From ad0dcee4df92d9384eeb12c98acc08683b04510d Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 1 Apr 2022 09:09:19 +0800
Subject: [PATCH 001/124] Initial commit.

---
 include/xgboost/objective.h     |  6 +++
 include/xgboost/tree_model.h    | 19 ++++++++
 include/xgboost/tree_updater.h  |  7 +++
 src/common/row_set.h            | 13 ++----
 src/gbm/gbtree.cc               |  4 +-
 src/objective/regression_obj.cc | 78 +++++++++++++++++++++++++++++++++
 src/tree/updater_approx.cc      | 41 +++++++++++++++--
 7 files changed, 153 insertions(+), 15 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 44dc46ddc8da..181a19393a13 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -22,6 +22,9 @@
 
 namespace xgboost {
 
+struct RowIndexCache;
+class RegTree;
+
 /*! \brief interface of objective function */
 class ObjFunction : public Configurable {
  protected:
@@ -88,6 +91,9 @@ class ObjFunction : public Configurable {
     return 1;
   }
 
+  virtual void UpdateTreeLeaf(RowIndexCache const& row_index, MetaInfo const& info, uint32_t target,
+                              RegTree* p_tree);
+
   /*!
    * \brief Create an objective function according to name.
    * \param tparam Generic parameters.
diff --git a/include/xgboost/tree_model.h b/include/xgboost/tree_model.h
index b2d2ad3383de..8079b86d5c94 100644
--- a/include/xgboost/tree_model.h
+++ b/include/xgboost/tree_model.h
@@ -734,5 +734,24 @@ inline bool RegTree::FVec::IsMissing(size_t i) const {
 inline bool RegTree::FVec::HasMissing() const {
   return has_missing_;
 }
+
+struct RowIndexCache {
+  struct Segment {
+    size_t begin;
+    size_t n;
+    bst_node_t nidx;
+  };
+
+  HostDeviceVector<size_t> row_index;
+  std::vector<Segment> indptr;
+
+  RowIndexCache(Context const* ctx, size_t n_leaf, size_t n_samples) {
+    indptr.resize(n_leaf + 1);
+    if (!ctx->IsCPU()) {
+      row_index.SetDevice(ctx->gpu_id);
+    }
+    row_index.Resize(n_samples);
+  }
+};
 }  // namespace xgboost
 #endif  // XGBOOST_TREE_MODEL_H_
diff --git a/include/xgboost/tree_updater.h b/include/xgboost/tree_updater.h
index 6189221dc0bf..f656d44bdafb 100644
--- a/include/xgboost/tree_updater.h
+++ b/include/xgboost/tree_updater.h
@@ -78,6 +78,13 @@ class TreeUpdater : public Configurable {
   }
 
   virtual char const* Name() const = 0;
+  /*!
+   * \brief Get the partition of rows based on the last iteration.
+   */
+  virtual common::Span<RowIndexCache> GetRowIndexCache(size_t tree_idx) {
+    LOG(FATAL) << "Objective is not supported by current tree method:" << this->Name();
+    return {};
+  }
 
   /*!
    * \brief Create a tree updater given name
diff --git a/src/common/row_set.h b/src/common/row_set.h
index dc61d5f5d877..7261b02cc714 100644
--- a/src/common/row_set.h
+++ b/src/common/row_set.h
@@ -15,7 +15,6 @@
 
 namespace xgboost {
 namespace common {
-
 /*! \brief collection of rowset */
 class RowSetCollection {
  public:
@@ -38,20 +37,17 @@ class RowSetCollection {
       return end - begin;
     }
   };
-  /* \brief specifies how to split a rowset into two */
-  struct Split {
-    std::vector<size_t> left;
-    std::vector<size_t> right;
-  };
 
-  inline std::vector<Elem>::const_iterator begin() const {  // NOLINT
+  std::vector<Elem>::const_iterator begin() const {  // NOLINT
     return elem_of_each_node_.begin();
   }
 
-  inline std::vector<Elem>::const_iterator end() const {  // NOLINT
+  std::vector<Elem>::const_iterator end() const {  // NOLINT
     return elem_of_each_node_.end();
   }
 
+  size_t Size() const { return std::distance(begin(), end()); }
+
   /*! \brief return corresponding element set given the node_id */
   inline const Elem& operator[](unsigned node_id) const {
     const Elem& e = elem_of_each_node_[node_id];
@@ -123,7 +119,6 @@ class RowSetCollection {
   // vector: node_id -> elements
   std::vector<Elem> elem_of_each_node_;
 };
-
 }  // namespace common
 }  // namespace xgboost
 
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index ec611ee95a68..450185d6488d 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -232,8 +232,7 @@ void GBTree::DoBoost(DMatrix* p_fmat,
   auto out = linalg::TensorView<float, 2>{
       device == GenericParameter::kCpuId ? predt->predictions.HostSpan()
                                          : predt->predictions.DeviceSpan(),
-      {static_cast<size_t>(p_fmat->Info().num_row_),
-       static_cast<size_t>(ngroup)},
+      {static_cast<size_t>(p_fmat->Info().num_row_), static_cast<size_t>(ngroup)},
       device};
   CHECK_NE(ngroup, 0);
   if (ngroup == 1) {
@@ -271,6 +270,7 @@ void GBTree::DoBoost(DMatrix* p_fmat,
       predt->Update(1);
     }
   }
+
   monitor_.Stop("BoostNewTrees");
   this->CommitModel(std::move(new_trees), p_fmat, predt);
 }
diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index 663989fbd5c3..ca44a1325e27 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -5,11 +5,89 @@
 // Dummy file to keep the CUDA conditional compile trick.
 
 #include <dmlc/registry.h>
+
+#include "../common/linalg_op.h"
+#include "rabit/rabit.h"
+#include "xgboost/data.h"
+#include "xgboost/objective.h"
+#include "xgboost/tree_model.h"
+
 namespace xgboost {
 namespace obj {
 
 DMLC_REGISTRY_FILE_TAG(regression_obj);
 
+float WeightedQuantile(float quantile, common::Span<size_t const> row_set,
+                       linalg::VectorView<float const> labels,
+                       linalg::VectorView<float const> weights) {
+  float result;
+  // fixme: pick an algorithm from R quantile.
+  return result;
+};
+
+float Quantile(float quantile, common::Span<size_t const> row_set, linalg::VectorView<float const> labels) {
+  float result;
+  // fixme: pick an algorithm from R quantile.
+  return result;
+}
+
+class MeanAbsoluteError : public ObjFunction {
+ public:
+  void Configure(Args const&) override {}
+
+  uint32_t Targets(MetaInfo const& info) const override {
+    return std::max(static_cast<size_t>(1), info.labels.Shape(1));
+  }
+
+  void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
+                   HostDeviceVector<GradientPair>* out_gpair) override {
+    out_gpair->SetDevice(ctx_->gpu_id);
+    out_gpair->Resize(info.labels.Size());
+    auto gpair = linalg::MakeVec(out_gpair);
+
+    preds.SetDevice(ctx_->gpu_id);
+    auto predt = linalg::MakeVec(&preds);
+
+    linalg::ElementWiseKernel(ctx_, info.labels.View(ctx_->gpu_id),
+                              [=] XGBOOST_DEVICE(size_t i, float const y) {});
+  }
+
+  void UpdateTreeLeaf(RowIndexCache const& row_index, MetaInfo const& info, uint32_t target,
+                      RegTree* p_tree) override {
+    auto& tree = *p_tree;
+    std::vector<float> results;
+    for (auto const& seg : row_index.indptr) {
+      auto h_row_set = row_index.row_index.HostSpan().subspan(seg.begin, seg.n);
+      float q{0};
+      if (info.weights_.Empty()) {
+        q = Quantile(0.5f, h_row_set, info.labels.HostView().Slice(linalg::All(), target));
+      } else {
+        q = WeightedQuantile(0.5f, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
+                             linalg::MakeVec(&info.weights_));
+      }
+      results.push_back(q);
+    }
+    // use the mean value
+    rabit::Allreduce<rabit::op::Sum>(results.data(), results.size());
+    auto world = rabit::GetWorldSize();
+    std::transform(results.begin(), results.end(), results.begin(),
+                   [&](float q) { return q / world; });
+    for (size_t i = 0; i < row_index.indptr.size(); ++i) {
+      auto seg = row_index.indptr[i];
+      auto q = results[i];
+      tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
+    }
+  }
+
+  const char* DefaultEvalMetric() const override { return "mae"; }
+
+  void SaveConfig(Json* p_out) const override {
+    auto& out = *p_out;
+    out["name"] = String("reg:mae");
+  }
+
+  void LoadConfig(Json const& in) override {}
+};
 }  // namespace obj
 }  // namespace xgboost
 
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index 3bad6f7da4cc..f13494db7fa6 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -164,8 +164,8 @@ class GloablApproxBuilder {
         ctx_{ctx},
         monitor_{monitor} {}
 
-  void UpdateTree(RegTree *p_tree, std::vector<GradientPair> const &gpair, common::Span<float> hess,
-                  DMatrix *p_fmat) {
+  void UpdateTree(DMatrix *p_fmat, std::vector<GradientPair> const &gpair, common::Span<float> hess,
+                  RegTree *p_tree, std::vector<RowIndexCache> *p_out_row_indices) {
     p_last_tree_ = p_tree;
     this->InitData(p_fmat, hess);
 
@@ -231,6 +231,29 @@ class GloablApproxBuilder {
       driver.Push(best_splits.begin(), best_splits.end());
       expand_set = driver.Pop();
     }
+
+    CHECK(p_out_row_indices->empty());
+    for (auto const &part : partitioner_) {
+      p_out_row_indices->emplace_back();
+      auto row_set = part.Partitions();
+      auto n_leaf = row_set.Size();
+      // fixme: subsample
+      auto &h_row_index = p_out_row_indices->back().row_index.HostVector();
+
+      auto begin = row_set.Data()->data();
+      for (auto node : row_set) {
+        CHECK(node.begin);
+        CHECK(tree[node.node_id].IsLeaf());
+        size_t offset = node.begin - begin;
+        auto size = node.Size();
+        auto seg = RowIndexCache::Segment{offset, size, node.node_id};
+        size_t k = seg.begin;
+        for (auto idx = node.begin; idx != node.end; ++idx) {
+          h_row_index[k] = *idx;
+        }
+        p_out_row_indices->back().indptr.push_back(seg);
+      }
+    }
   }
 };
 
@@ -249,6 +272,8 @@ class GlobalApproxUpdater : public TreeUpdater {
   DMatrix *cached_{nullptr};
   std::shared_ptr<common::ColumnSampler> column_sampler_ =
       std::make_shared<common::ColumnSampler>();
+  // cache for row partitions
+  std::vector<std::vector<RowIndexCache>> row_set_collection_;
   ObjInfo task_;
 
  public:
@@ -275,6 +300,8 @@ class GlobalApproxUpdater : public TreeUpdater {
     sampled->resize(h_gpair.size());
     std::copy(h_gpair.cbegin(), h_gpair.cend(), sampled->begin());
     auto &rnd = common::GlobalRandom();
+    row_set_collection_.clear();
+
     if (param.subsample != 1.0) {
       CHECK(param.sampling_method != TrainParam::kGradientBased)
           << "Gradient based sampling is not supported for approx tree method.";
@@ -314,10 +341,12 @@ class GlobalApproxUpdater : public TreeUpdater {
     cached_ = m;
 
     for (auto p_tree : trees) {
+      row_set_collection_.emplace_back();
+      auto &row_indices = row_set_collection_.back();
       if (hist_param_.single_precision_histogram) {
-        this->f32_impl_->UpdateTree(p_tree, h_gpair, hess, m);
+        this->f32_impl_->UpdateTree(m, h_gpair, hess, p_tree, &row_indices);
       } else {
-        this->f64_impl_->UpdateTree(p_tree, h_gpair, hess, m);
+        this->f64_impl_->UpdateTree(m, h_gpair, hess, p_tree, &row_indices);
       }
     }
     param_.learning_rate = lr;
@@ -335,6 +364,10 @@ class GlobalApproxUpdater : public TreeUpdater {
     }
     return true;
   }
+
+  common::Span<RowIndexCache> GetRowIndexCache(size_t tree_idx) override {
+    return row_set_collection_.at(tree_idx);
+  }
 };
 
 DMLC_REGISTRY_FILE_TAG(grow_histmaker);

From 73aef42c1aa328bd0aba1d95efffd265c8415b4b Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 1 Apr 2022 10:57:34 +0800
Subject: [PATCH 002/124] refresh.

---
 include/xgboost/gbm.h           |   5 +-
 include/xgboost/objective.h     |   4 +-
 include/xgboost/tree_updater.h  |   2 +-
 src/common/common.h             |   2 +-
 src/gbm/gblinear.cc             |   5 +-
 src/gbm/gbtree.cc               |  42 +++++++------
 src/gbm/gbtree.h                |   5 +-
 src/objective/regression_obj.cc | 103 ++++++++++++++++++++++++--------
 src/tree/updater_approx.cc      |   2 +-
 9 files changed, 115 insertions(+), 55 deletions(-)

diff --git a/include/xgboost/gbm.h b/include/xgboost/gbm.h
index d24057e255a7..cce92d3679f4 100644
--- a/include/xgboost/gbm.h
+++ b/include/xgboost/gbm.h
@@ -90,9 +90,8 @@ class GradientBooster : public Model, public Configurable {
    * \param prediction The output prediction cache entry that needs to be updated.
    * the booster may change content of gpair
    */
-  virtual void DoBoost(DMatrix* p_fmat,
-                       HostDeviceVector<GradientPair>* in_gpair,
-                       PredictionCacheEntry*) = 0;
+  virtual void DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
+                       PredictionCacheEntry*, ObjFunction const* obj) = 0;
 
   /*!
    * \brief generate predictions for given feature matrix
diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 181a19393a13..a0ead2ff11be 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -91,8 +91,8 @@ class ObjFunction : public Configurable {
     return 1;
   }
 
-  virtual void UpdateTreeLeaf(RowIndexCache const& row_index, MetaInfo const& info, uint32_t target,
-                              RegTree* p_tree);
+  virtual void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
+                              uint32_t target, RegTree* p_tree) const;
 
   /*!
    * \brief Create an objective function according to name.
diff --git a/include/xgboost/tree_updater.h b/include/xgboost/tree_updater.h
index f656d44bdafb..94264e3b6b8e 100644
--- a/include/xgboost/tree_updater.h
+++ b/include/xgboost/tree_updater.h
@@ -81,7 +81,7 @@ class TreeUpdater : public Configurable {
   /*!
    * \brief Get the partition of rows based on the last iteration.
    */
-  virtual common::Span<RowIndexCache> GetRowIndexCache(size_t tree_idx) {
+  virtual common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const {
     LOG(FATAL) << "Objective is not supported by current tree method:" << this->Name();
     return {};
   }
diff --git a/src/common/common.h b/src/common/common.h
index fb7e7fee55da..066897e44b45 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -191,7 +191,7 @@ std::vector<Idx> ArgSort(Container const &array, Comp comp = std::less<V>{}) {
 
 struct OptionalWeights {
   Span<float const> weights;
-  float dft{1.0f};
+  float dft{1.0f};  // fixme: make this compile time constant
 
   explicit OptionalWeights(Span<float const> w) : weights{w} {}
   explicit OptionalWeights(float w) : dft{w} {}
diff --git a/src/gbm/gblinear.cc b/src/gbm/gblinear.cc
index cbf6ffebfca5..0e983fe4b37f 100644
--- a/src/gbm/gblinear.cc
+++ b/src/gbm/gblinear.cc
@@ -134,9 +134,8 @@ class GBLinear : public GradientBooster {
     this->updater_->SaveConfig(&j_updater);
   }
 
-  void DoBoost(DMatrix *p_fmat,
-               HostDeviceVector<GradientPair> *in_gpair,
-               PredictionCacheEntry*) override {
+  void DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair, PredictionCacheEntry*,
+               ObjFunction const*) override {
     monitor_.Start("DoBoost");
 
     model_.LazyInitModel();
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index 450185d6488d..372559d41205 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -1,33 +1,34 @@
 /*!
- * Copyright 2014-2021 by Contributors
+ * Copyright 2014-2022 by Contributors
  * \file gbtree.cc
  * \brief gradient boosted tree implementation.
  * \author Tianqi Chen
  */
+#include "gbtree.h"
+
 #include <dmlc/omp.h>
 #include <dmlc/parameter.h>
 
-#include <vector>
+#include <algorithm>
+#include <limits>
 #include <memory>
-#include <utility>
 #include <string>
-#include <limits>
-#include <algorithm>
+#include <utility>
+#include <vector>
 
+#include "../common/common.h"
+#include "../common/random.h"
+#include "../common/threading_utils.h"
+#include "../common/timer.h"
+#include "gbtree_model.h"
 #include "xgboost/data.h"
 #include "xgboost/gbm.h"
-#include "xgboost/logging.h"
+#include "xgboost/host_device_vector.h"
 #include "xgboost/json.h"
+#include "xgboost/logging.h"
+#include "xgboost/objective.h"
 #include "xgboost/predictor.h"
 #include "xgboost/tree_updater.h"
-#include "xgboost/host_device_vector.h"
-
-#include "gbtree.h"
-#include "gbtree_model.h"
-#include "../common/common.h"
-#include "../common/random.h"
-#include "../common/timer.h"
-#include "../common/threading_utils.h"
 
 namespace xgboost {
 namespace gbm {
@@ -216,9 +217,8 @@ void CopyGradient(HostDeviceVector<GradientPair> const* in_gpair, int32_t n_thre
   }
 }
 
-void GBTree::DoBoost(DMatrix* p_fmat,
-                     HostDeviceVector<GradientPair>* in_gpair,
-                     PredictionCacheEntry* predt) {
+void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
+                     PredictionCacheEntry* predt, ObjFunction const* obj) {
   std::vector<std::vector<std::unique_ptr<RegTree> > > new_trees;
   const int ngroup = model_.learner_model_param->num_output_group;
   ConfigureWithKnownData(this->cfg_, p_fmat);
@@ -271,6 +271,14 @@ void GBTree::DoBoost(DMatrix* p_fmat,
     }
   }
 
+  bst_group_t gidx {0};
+  for (auto& tree_group : new_trees) {
+    for (size_t t = 0; t < tree_group.size(); ++t) {
+      auto row_idx = updaters_.back()->GetRowIndexCache(t);
+      auto target = p_fmat->Info().labels.Shape(1) > 1 ? gidx : 0;
+      obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), gidx, tree_group[t].get());
+    }
+  }
   monitor_.Stop("BoostNewTrees");
   this->CommitModel(std::move(new_trees), p_fmat, predt);
 }
diff --git a/src/gbm/gbtree.h b/src/gbm/gbtree.h
index 67d9e212888a..c973d0b02685 100644
--- a/src/gbm/gbtree.h
+++ b/src/gbm/gbtree.h
@@ -203,9 +203,8 @@ class GBTree : public GradientBooster {
   void ConfigureWithKnownData(Args const& cfg, DMatrix* fmat);
 
   /*! \brief Carry out one iteration of boosting */
-  void DoBoost(DMatrix* p_fmat,
-               HostDeviceVector<GradientPair>* in_gpair,
-               PredictionCacheEntry* predt) override;
+  void DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
+               PredictionCacheEntry* predt, ObjFunction const* obj) override;
 
   bool UseGPU() const override {
     return
diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index ca44a1325e27..85052f78853d 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -20,13 +20,37 @@ DMLC_REGISTRY_FILE_TAG(regression_obj);
 float WeightedQuantile(float quantile, common::Span<size_t const> row_set,
                        linalg::VectorView<float const> labels,
                        linalg::VectorView<float const> weights) {
-  float result;
-  // fixme: pick an algorithm from R quantile.
-  return result;
+  std::vector<size_t> sorted_idx(row_set.size());
+  std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
+  std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
+                   [&](size_t i, size_t j) { return labels(row_set[i]) < labels(row_set[j]); });
+  std::vector<float> weighted_cdf(row_set.size());
+  weighted_cdf[0] = weights(row_set[sorted_idx[0]]);
+  for (size_t i = 1; i < row_set.size(); ++i) {
+    weighted_cdf[i] = weighted_cdf[i - 1] + weights(row_set[sorted_idx[i]]);
+  }
+  float thresh = weighted_cdf.back() * quantile;
+  size_t pos =
+      std::upper_bound(weighted_cdf.cbegin(), weighted_cdf.cend(), thresh) - weighted_cdf.cbegin();
+  pos = std::min(pos, static_cast<size_t>(row_set.size() - 1));
+  if (pos == 0 || pos == static_cast<size_t>(row_set.size() - 1)) {
+    return labels(row_set[sorted_idx[pos]]);
+  }
+  CHECK_GE(thresh, weighted_cdf[pos - 1]);
+  CHECK_LT(thresh, weighted_cdf[pos]);
+  float v1 = labels(row_set[sorted_idx[pos - 1]]);
+  float v2 = labels(row_set[sorted_idx[pos]]);
+  if (weighted_cdf[pos + 1] - weighted_cdf[pos] >= 1.0f) {
+    return (thresh - weighted_cdf[pos]) / (weighted_cdf[pos + 1] - weighted_cdf[pos]) * (v2 - v2) +
+           v1;
+  } else {
+    return v2;
+  }
 };
 
 float Quantile(float quantile, common::Span<size_t const> row_set, linalg::VectorView<float const> labels) {
   float result;
+  LOG(FATAL) << "Not implemented";
   // fixme: pick an algorithm from R quantile.
   return result;
 }
@@ -39,6 +63,10 @@ class MeanAbsoluteError : public ObjFunction {
     return std::max(static_cast<size_t>(1), info.labels.Shape(1));
   }
 
+  struct ObjInfo Task() const override {
+    return {ObjInfo::kRegression, true};
+  }
+
   void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
                    HostDeviceVector<GradientPair>* out_gpair) override {
     out_gpair->SetDevice(ctx_->gpu_id);
@@ -47,34 +75,57 @@ class MeanAbsoluteError : public ObjFunction {
 
     preds.SetDevice(ctx_->gpu_id);
     auto predt = linalg::MakeVec(&preds);
+    auto sign = [](auto x) {
+      return (x > static_cast<decltype(x)>(0)) - (x < static_cast<decltype(x)>(0));
+    };
+
+    info.weights_.SetDevice(ctx_->gpu_id);
+    common::OptionalWeights weight{ctx_->IsCPU() ? info.weights_.ConstHostSpan()
+                                                 : info.weights_.ConstDeviceSpan()};
 
     linalg::ElementWiseKernel(ctx_, info.labels.View(ctx_->gpu_id),
-                              [=] XGBOOST_DEVICE(size_t i, float const y) {});
+                              [=] XGBOOST_DEVICE(size_t i, float const y) mutable {
+                                auto grad = sign(predt(i) - y) * weight[i];
+                                auto hess = weight[i];
+                                gpair(i) = GradientPair{grad, hess};
+                              });
   }
 
-  void UpdateTreeLeaf(RowIndexCache const& row_index, MetaInfo const& info, uint32_t target,
-                      RegTree* p_tree) override {
+  void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
+                      uint32_t target, RegTree* p_tree) const override {
     auto& tree = *p_tree;
-    std::vector<float> results;
-    for (auto const& seg : row_index.indptr) {
-      auto h_row_set = row_index.row_index.HostSpan().subspan(seg.begin, seg.n);
-      float q{0};
-      if (info.weights_.Empty()) {
-        q = Quantile(0.5f, h_row_set, info.labels.HostView().Slice(linalg::All(), target));
-      } else {
-        q = WeightedQuantile(0.5f, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
-                             linalg::MakeVec(&info.weights_));
+    std::vector<float> quantiles;
+    for (auto const& part : row_index) {
+      std::vector<float> results;
+      for (auto const& seg : part.indptr) {
+        auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
+        float q{0};
+        if (info.weights_.Empty()) {
+          q = Quantile(0.5f, h_row_set, info.labels.HostView().Slice(linalg::All(), target));
+        } else {
+          q = WeightedQuantile(0.5f, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
+                               linalg::MakeVec(&info.weights_));
+        }
+        results.push_back(q);
+      }
+      // fixme: verify this is correct for external memory
+      if (quantiles.empty()) {
+        quantiles.resize(results.size(), 0);
       }
-      results.push_back(q);
+      for (size_t i = 0; i < results.size(); ++i) {
+        quantiles[i] += results[i];
+      }
+      // use the mean value
+      rabit::Allreduce<rabit::op::Sum>(results.data(), results.size());
+      auto world = rabit::GetWorldSize();
+      std::transform(results.begin(), results.end(), results.begin(),
+                     [&](float q) { return q / world; });
     }
-    // use the mean value
-    rabit::Allreduce<rabit::op::Sum>(results.data(), results.size());
-    auto world = rabit::GetWorldSize();
-    std::transform(results.begin(), results.end(), results.begin(),
-                   [&](float q) { return q / world; });
-    for (size_t i = 0; i < row_index.indptr.size(); ++i) {
-      auto seg = row_index.indptr[i];
-      auto q = results[i];
+
+    // fixme: verify this is correct for external memory
+    for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
+      auto seg = row_index.front().indptr[i];
+      auto q = quantiles[i];
       tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
     }
   }
@@ -88,6 +139,10 @@ class MeanAbsoluteError : public ObjFunction {
 
   void LoadConfig(Json const& in) override {}
 };
+
+XGBOOST_REGISTER_OBJECTIVE(MeanAbsoluteError, "reg:absoluteerror")
+    .describe("Regression Pseudo Huber error.")
+    .set_body([]() { return new MeanAbsoluteError(); });
 }  // namespace obj
 }  // namespace xgboost
 
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index f13494db7fa6..0aa26cf7753e 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -365,7 +365,7 @@ class GlobalApproxUpdater : public TreeUpdater {
     return true;
   }
 
-  common::Span<RowIndexCache> GetRowIndexCache(size_t tree_idx) override {
+  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
     return row_set_collection_.at(tree_idx);
   }
 };

From f0d949a0a65bfb405947c0c7ec770993b3a98a0d Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 1 Apr 2022 11:08:50 +0800
Subject: [PATCH 003/124] Start looking into quantile reg.

---
 src/objective/regression_obj.cc | 119 ++++++++++++++++++++++----------
 1 file changed, 82 insertions(+), 37 deletions(-)

diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index 85052f78853d..caf6d1f8d6ba 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -55,6 +55,45 @@ float Quantile(float quantile, common::Span<size_t const> row_set, linalg::Vecto
   return result;
 }
 
+void UpdateTreeLeafHost(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
+                        uint32_t target, float alpha, RegTree* p_tree) {
+  auto& tree = *p_tree;
+  std::vector<float> quantiles;
+  for (auto const& part : row_index) {
+    std::vector<float> results;
+    for (auto const& seg : part.indptr) {
+      auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
+      float q{0};
+      if (info.weights_.Empty()) {
+        q = Quantile(alpha, h_row_set, info.labels.HostView().Slice(linalg::All(), target));
+      } else {
+        q = WeightedQuantile(alpha, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
+                             linalg::MakeVec(&info.weights_));
+      }
+      results.push_back(q);
+    }
+    // fixme: verify this is correct for external memory
+    if (quantiles.empty()) {
+      quantiles.resize(results.size(), 0);
+    }
+    for (size_t i = 0; i < results.size(); ++i) {
+      quantiles[i] += results[i];
+    }
+    // use the mean value
+    rabit::Allreduce<rabit::op::Sum>(results.data(), results.size());
+    auto world = rabit::GetWorldSize();
+    std::transform(results.begin(), results.end(), results.begin(),
+                   [&](float q) { return q / world; });
+  }
+
+  // fixme: verify this is correct for external memory
+  for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
+    auto seg = row_index.front().indptr[i];
+    auto q = quantiles[i];
+    tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
+  }
+}
+
 class MeanAbsoluteError : public ObjFunction {
  public:
   void Configure(Args const&) override {}
@@ -93,56 +132,62 @@ class MeanAbsoluteError : public ObjFunction {
 
   void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
                       uint32_t target, RegTree* p_tree) const override {
-    auto& tree = *p_tree;
-    std::vector<float> quantiles;
-    for (auto const& part : row_index) {
-      std::vector<float> results;
-      for (auto const& seg : part.indptr) {
-        auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
-        float q{0};
-        if (info.weights_.Empty()) {
-          q = Quantile(0.5f, h_row_set, info.labels.HostView().Slice(linalg::All(), target));
-        } else {
-          q = WeightedQuantile(0.5f, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
-                               linalg::MakeVec(&info.weights_));
-        }
-        results.push_back(q);
-      }
-      // fixme: verify this is correct for external memory
-      if (quantiles.empty()) {
-        quantiles.resize(results.size(), 0);
-      }
-      for (size_t i = 0; i < results.size(); ++i) {
-        quantiles[i] += results[i];
-      }
-      // use the mean value
-      rabit::Allreduce<rabit::op::Sum>(results.data(), results.size());
-      auto world = rabit::GetWorldSize();
-      std::transform(results.begin(), results.end(), results.begin(),
-                     [&](float q) { return q / world; });
-    }
-
-    // fixme: verify this is correct for external memory
-    for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
-      auto seg = row_index.front().indptr[i];
-      auto q = quantiles[i];
-      tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
-    }
+    UpdateTreeLeafHost(row_index, info, target, 0.5, p_tree);
   }
 
   const char* DefaultEvalMetric() const override { return "mae"; }
 
   void SaveConfig(Json* p_out) const override {
     auto& out = *p_out;
-    out["name"] = String("reg:mae");
+    out["name"] = String("reg:absoluteerror");
   }
 
   void LoadConfig(Json const& in) override {}
 };
 
 XGBOOST_REGISTER_OBJECTIVE(MeanAbsoluteError, "reg:absoluteerror")
-    .describe("Regression Pseudo Huber error.")
+    .describe("Mean absoluate error.")
     .set_body([]() { return new MeanAbsoluteError(); });
+
+struct QuantileRegressionParameter : public XGBoostParameter<QuantileRegressionParameter> {
+  float quantile;
+};
+
+class QuantileRegression : public ObjFunction {
+  QuantileRegressionParameter param_;
+
+ public:
+  void Configure(Args const&) override {}
+
+  uint32_t Targets(MetaInfo const& info) const override {
+    return std::max(static_cast<size_t>(1), info.labels.Shape(1));
+  }
+
+  struct ObjInfo Task() const override {
+    return {ObjInfo::kRegression, true};
+  }
+
+  void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
+                   HostDeviceVector<GradientPair>* out_gpair) override {}
+
+  void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
+                      uint32_t target, RegTree* p_tree) const override {
+    UpdateTreeLeafHost(row_index, info, target, param_.quantile, p_tree);
+  }
+
+  const char* DefaultEvalMetric() const override { return "undefined"; }
+
+  void SaveConfig(Json* p_out) const override {
+    auto& out = *p_out;
+    out["name"] = String("reg:quantile");
+    out["quantile_regression_param"] = ToJson(param_);
+  }
+  void LoadConfig(Json const& in) override { FromJson(in["quantile_regression_param"], &param_); }
+};
+
+XGBOOST_REGISTER_OBJECTIVE(QuantileRegression, "reg:quantile")
+    .describe("Quantile regression.")
+    .set_body([]() { return new QuantileRegression(); });
 }  // namespace obj
 }  // namespace xgboost
 

From 3def77bdaced8c26a0980a99e2f8dc02e723a905 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 1 Apr 2022 11:19:07 +0800
Subject: [PATCH 004/124] init.

---
 src/objective/regression_obj.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index caf6d1f8d6ba..333bf8c30d32 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -150,7 +150,7 @@ XGBOOST_REGISTER_OBJECTIVE(MeanAbsoluteError, "reg:absoluteerror")
     .set_body([]() { return new MeanAbsoluteError(); });
 
 struct QuantileRegressionParameter : public XGBoostParameter<QuantileRegressionParameter> {
-  float quantile;
+  float quantile{0.5};
 };
 
 class QuantileRegression : public ObjFunction {

From 4f5cd8c2e3656c57ad36d00b56e4f4d8295fd9f4 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 1 Apr 2022 11:27:11 +0800
Subject: [PATCH 005/124] quantile.

---
 src/objective/regression_obj.cc | 39 +++++++++++++++++++++++++++------
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index 333bf8c30d32..1d6de1575e92 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -108,6 +108,8 @@ class MeanAbsoluteError : public ObjFunction {
 
   void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
                    HostDeviceVector<GradientPair>* out_gpair) override {
+    auto labels = info.labels.View(ctx_->gpu_id);
+
     out_gpair->SetDevice(ctx_->gpu_id);
     out_gpair->Resize(info.labels.Size());
     auto gpair = linalg::MakeVec(out_gpair);
@@ -122,12 +124,12 @@ class MeanAbsoluteError : public ObjFunction {
     common::OptionalWeights weight{ctx_->IsCPU() ? info.weights_.ConstHostSpan()
                                                  : info.weights_.ConstDeviceSpan()};
 
-    linalg::ElementWiseKernel(ctx_, info.labels.View(ctx_->gpu_id),
-                              [=] XGBOOST_DEVICE(size_t i, float const y) mutable {
-                                auto grad = sign(predt(i) - y) * weight[i];
-                                auto hess = weight[i];
-                                gpair(i) = GradientPair{grad, hess};
-                              });
+    linalg::ElementWiseKernel(ctx_, labels, [=] XGBOOST_DEVICE(size_t i, float const y) mutable {
+      auto sample_id = std::get<0>(linalg::UnravelIndex(i, labels.Shape()));
+      auto grad = sign(predt(i) - y) * weight[i];
+      auto hess = weight[sample_id];
+      gpair(i) = GradientPair{grad, hess};
+    });
   }
 
   void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
@@ -168,7 +170,30 @@ class QuantileRegression : public ObjFunction {
   }
 
   void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
-                   HostDeviceVector<GradientPair>* out_gpair) override {}
+                   HostDeviceVector<GradientPair>* out_gpair) override {
+    auto labels = info.labels.View(ctx_->gpu_id);
+
+    out_gpair->SetDevice(ctx_->gpu_id);
+    out_gpair->Resize(info.labels.Size());
+    auto gpair = linalg::MakeVec(out_gpair);
+
+    preds.SetDevice(ctx_->gpu_id);
+    auto predt = linalg::MakeVec(&preds);
+    auto quantile = param_.quantile;
+
+    info.weights_.SetDevice(ctx_->gpu_id);
+    common::OptionalWeights weight{ctx_->IsCPU() ? info.weights_.ConstHostSpan()
+                                                 : info.weights_.ConstDeviceSpan()};
+
+    linalg::ElementWiseKernel(ctx_, labels, [=] XGBOOST_DEVICE(size_t i, float const y) mutable {
+      auto sample_id = std::get<0>(linalg::UnravelIndex(i, labels.Shape()));
+      auto res = predt(i) - y;
+      auto grad = res >= 0 ? (1.0f - quantile) : -quantile;
+      grad *= weight[sample_id];
+      auto hess = weight[sample_id];
+      gpair(i) = GradientPair{grad, hess};
+    });
+  }
 
   void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
                       uint32_t target, RegTree* p_tree) const override {

From b216a38a3ac2071dbaa9fd82983d982083662a00 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Thu, 7 Apr 2022 18:54:09 +0800
Subject: [PATCH 006/124] percentile.

---
 include/xgboost/objective.h               |   2 +-
 src/common/common.h                       |  18 +-
 src/common/stats.h                        |  31 ++
 src/learner.cc                            |   4 +-
 src/objective/regression_obj.cc           |  63 ----
 src/tree/updater_approx.cc                | 383 +---------------------
 tests/cpp/common/test_stats.cc            |  14 +
 tests/cpp/gbm/test_gbtree.cc              |   6 +-
 tests/cpp/helpers.cc                      |   2 +-
 tests/cpp/predictor/test_cpu_predictor.cc |   2 +-
 10 files changed, 69 insertions(+), 456 deletions(-)
 create mode 100644 src/common/stats.h
 create mode 100644 tests/cpp/common/test_stats.cc

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index a0ead2ff11be..385889d4784b 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -92,7 +92,7 @@ class ObjFunction : public Configurable {
   }
 
   virtual void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
-                              uint32_t target, RegTree* p_tree) const;
+                              uint32_t target, RegTree* p_tree) const {}
 
   /*!
    * \brief Create an objective function according to name.
diff --git a/src/common/common.h b/src/common/common.h
index 066897e44b45..877feba81553 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -7,6 +7,7 @@
 #define XGBOOST_COMMON_COMMON_H_
 
 #include <xgboost/base.h>
+#include <xgboost/linalg.h>
 #include <xgboost/logging.h>
 #include <xgboost/span.h>
 
@@ -14,12 +15,12 @@
 #include <exception>
 #include <functional>
 #include <limits>
-#include <type_traits>
-#include <vector>
-#include <string>
-#include <sstream>
 #include <numeric>
+#include <sstream>
+#include <string>
+#include <type_traits>
 #include <utility>
+#include <vector>
 
 #if defined(__CUDACC__)
 #include <thrust/system/cuda/error.h>
@@ -189,6 +190,15 @@ std::vector<Idx> ArgSort(Container const &array, Comp comp = std::less<V>{}) {
   return result;
 }
 
+template <typename V, typename Comp = std::less<V>>
+std::vector<size_t> ArgSort(linalg::TensorView<V, 1> array, Comp comp = std::less<V>{}) {
+  std::vector<size_t> result(array.Size());
+  std::iota(result.begin(), result.end(), 0);
+  auto op = [&array, comp](size_t const &l, size_t const &r) { return comp(array(l), array(r)); };
+  XGBOOST_PARALLEL_STABLE_SORT(result.begin(), result.end(), op);
+  return result;
+}
+
 struct OptionalWeights {
   Span<float const> weights;
   float dft{1.0f};  // fixme: make this compile time constant
diff --git a/src/common/stats.h b/src/common/stats.h
new file mode 100644
index 000000000000..c6fae6bb0a7c
--- /dev/null
+++ b/src/common/stats.h
@@ -0,0 +1,31 @@
+#include <limits>
+#include <vector>
+
+#include "common.h"
+#include "xgboost/linalg.h"
+
+namespace xgboost {
+namespace common {
+float Percentile(float percentile, linalg::TensorView<float const, 1> arr) {
+  size_t n = arr.Shape(0);
+  if (n == 0) {
+    return std::numeric_limits<float>::quiet_NaN();
+  }
+  std::vector<size_t> sorted_idx{ArgSort(arr)};
+
+  if (percentile <= (1 / (n + 1))) {
+    return arr(sorted_idx.front());
+  }
+  if (percentile >= (n / (n + 1))) {
+    return arr(sorted_idx.back());
+  }
+  double x = percentile * static_cast<double>((n + 1));
+  double k = std::floor(x);
+  double d = x - k;
+
+  auto v0 = arr(sorted_idx[static_cast<size_t>(k)]);
+  auto v1 = arr(sorted_idx[static_cast<size_t>(k) + 1]);
+  return v0 + d * (v1 - v0);
+}
+}  // namespace common
+}  // namespace xgboost
diff --git a/src/learner.cc b/src/learner.cc
index 73447cf2ef1a..becc45553bbf 100644
--- a/src/learner.cc
+++ b/src/learner.cc
@@ -1168,7 +1168,7 @@ class LearnerImpl : public LearnerIO {
     monitor_.Stop("GetGradient");
     TrainingObserver::Instance().Observe(gpair_, "Gradients");
 
-    gbm_->DoBoost(train.get(), &gpair_, &predt);
+    gbm_->DoBoost(train.get(), &gpair_, &predt, obj_.get());
     monitor_.Stop("UpdateOneIter");
   }
 
@@ -1185,7 +1185,7 @@ class LearnerImpl : public LearnerIO {
     auto local_cache = this->GetPredictionCache();
     local_cache->Cache(train, generic_parameters_.gpu_id);
 
-    gbm_->DoBoost(train.get(), in_gpair, &local_cache->Entry(train.get()));
+    gbm_->DoBoost(train.get(), in_gpair, &local_cache->Entry(train.get()), obj_.get());
     monitor_.Stop("BoostOneIter");
   }
 
diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index 1d6de1575e92..75afd78963c4 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -150,69 +150,6 @@ class MeanAbsoluteError : public ObjFunction {
 XGBOOST_REGISTER_OBJECTIVE(MeanAbsoluteError, "reg:absoluteerror")
     .describe("Mean absoluate error.")
     .set_body([]() { return new MeanAbsoluteError(); });
-
-struct QuantileRegressionParameter : public XGBoostParameter<QuantileRegressionParameter> {
-  float quantile{0.5};
-};
-
-class QuantileRegression : public ObjFunction {
-  QuantileRegressionParameter param_;
-
- public:
-  void Configure(Args const&) override {}
-
-  uint32_t Targets(MetaInfo const& info) const override {
-    return std::max(static_cast<size_t>(1), info.labels.Shape(1));
-  }
-
-  struct ObjInfo Task() const override {
-    return {ObjInfo::kRegression, true};
-  }
-
-  void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
-                   HostDeviceVector<GradientPair>* out_gpair) override {
-    auto labels = info.labels.View(ctx_->gpu_id);
-
-    out_gpair->SetDevice(ctx_->gpu_id);
-    out_gpair->Resize(info.labels.Size());
-    auto gpair = linalg::MakeVec(out_gpair);
-
-    preds.SetDevice(ctx_->gpu_id);
-    auto predt = linalg::MakeVec(&preds);
-    auto quantile = param_.quantile;
-
-    info.weights_.SetDevice(ctx_->gpu_id);
-    common::OptionalWeights weight{ctx_->IsCPU() ? info.weights_.ConstHostSpan()
-                                                 : info.weights_.ConstDeviceSpan()};
-
-    linalg::ElementWiseKernel(ctx_, labels, [=] XGBOOST_DEVICE(size_t i, float const y) mutable {
-      auto sample_id = std::get<0>(linalg::UnravelIndex(i, labels.Shape()));
-      auto res = predt(i) - y;
-      auto grad = res >= 0 ? (1.0f - quantile) : -quantile;
-      grad *= weight[sample_id];
-      auto hess = weight[sample_id];
-      gpair(i) = GradientPair{grad, hess};
-    });
-  }
-
-  void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
-                      uint32_t target, RegTree* p_tree) const override {
-    UpdateTreeLeafHost(row_index, info, target, param_.quantile, p_tree);
-  }
-
-  const char* DefaultEvalMetric() const override { return "undefined"; }
-
-  void SaveConfig(Json* p_out) const override {
-    auto& out = *p_out;
-    out["name"] = String("reg:quantile");
-    out["quantile_regression_param"] = ToJson(param_);
-  }
-  void LoadConfig(Json const& in) override { FromJson(in["quantile_regression_param"], &param_); }
-};
-
-XGBOOST_REGISTER_OBJECTIVE(QuantileRegression, "reg:quantile")
-    .describe("Quantile regression.")
-    .set_body([]() { return new QuantileRegression(); });
 }  // namespace obj
 }  // namespace xgboost
 
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index 0aa26cf7753e..fd01034815bc 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -1,381 +1,2 @@
-/*!
- * Copyright 2021-2022 XGBoost contributors
- *
- * \brief Implementation for the approx tree method.
- */
-#include "updater_approx.h"
-
-#include <algorithm>
-#include <memory>
-#include <vector>
-
-#include "../common/random.h"
-#include "../data/gradient_index.h"
-#include "constraints.h"
-#include "driver.h"
-#include "hist/evaluate_splits.h"
-#include "hist/histogram.h"
-#include "hist/param.h"
-#include "param.h"
-#include "xgboost/base.h"
-#include "xgboost/json.h"
-#include "xgboost/tree_updater.h"
-
-namespace xgboost {
-namespace tree {
-
-DMLC_REGISTRY_FILE_TAG(updater_approx);
-
-namespace {
-// Return the BatchParam used by DMatrix.
-template <typename GradientSumT>
-auto BatchSpec(TrainParam const &p, common::Span<float> hess,
-               HistEvaluator<GradientSumT, CPUExpandEntry> const &evaluator) {
-  return BatchParam{p.max_bin, hess, !evaluator.Task().const_hess};
-}
-
-auto BatchSpec(TrainParam const &p, common::Span<float> hess) {
-  return BatchParam{p.max_bin, hess, false};
-}
-}  // anonymous namespace
-
-template <typename GradientSumT>
-class GloablApproxBuilder {
- protected:
-  TrainParam param_;
-  std::shared_ptr<common::ColumnSampler> col_sampler_;
-  HistEvaluator<GradientSumT, CPUExpandEntry> evaluator_;
-  HistogramBuilder<GradientSumT, CPUExpandEntry> histogram_builder_;
-  GenericParameter const *ctx_;
-
-  std::vector<ApproxRowPartitioner> partitioner_;
-  // Pointer to last updated tree, used for update prediction cache.
-  RegTree *p_last_tree_{nullptr};
-  common::Monitor *monitor_;
-  size_t n_batches_{0};
-  // Cache for histogram cuts.
-  common::HistogramCuts feature_values_;
-
- public:
-  void InitData(DMatrix *p_fmat, common::Span<float> hess) {
-    monitor_->Start(__func__);
-
-    n_batches_ = 0;
-    int32_t n_total_bins = 0;
-    partitioner_.clear();
-    // Generating the GHistIndexMatrix is quite slow, is there a way to speed it up?
-    for (auto const &page :
-         p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess, evaluator_))) {
-      if (n_total_bins == 0) {
-        n_total_bins = page.cut.TotalBins();
-        feature_values_ = page.cut;
-      } else {
-        CHECK_EQ(n_total_bins, page.cut.TotalBins());
-      }
-      partitioner_.emplace_back(page.Size(), page.base_rowid);
-      n_batches_++;
-    }
-
-    histogram_builder_.Reset(n_total_bins, BatchSpec(param_, hess), ctx_->Threads(), n_batches_,
-                             rabit::IsDistributed());
-    monitor_->Stop(__func__);
-  }
-
-  CPUExpandEntry InitRoot(DMatrix *p_fmat, std::vector<GradientPair> const &gpair,
-                          common::Span<float> hess, RegTree *p_tree) {
-    monitor_->Start(__func__);
-    CPUExpandEntry best;
-    best.nid = RegTree::kRoot;
-    best.depth = 0;
-    GradStats root_sum;
-    for (auto const &g : gpair) {
-      root_sum.Add(g);
-    }
-    rabit::Allreduce<rabit::op::Sum, double>(reinterpret_cast<double *>(&root_sum), 2);
-    std::vector<CPUExpandEntry> nodes{best};
-    size_t i = 0;
-    auto space = ConstructHistSpace(partitioner_, nodes);
-    for (auto const &page : p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess))) {
-      histogram_builder_.BuildHist(i, space, page, p_tree, partitioner_.at(i).Partitions(), nodes,
-                                   {}, gpair);
-      i++;
-    }
-
-    auto weight = evaluator_.InitRoot(root_sum);
-    p_tree->Stat(RegTree::kRoot).sum_hess = root_sum.GetHess();
-    p_tree->Stat(RegTree::kRoot).base_weight = weight;
-    (*p_tree)[RegTree::kRoot].SetLeaf(param_.learning_rate * weight);
-
-    auto const &histograms = histogram_builder_.Histogram();
-    auto ft = p_fmat->Info().feature_types.ConstHostSpan();
-    evaluator_.EvaluateSplits(histograms, feature_values_, ft, *p_tree, &nodes);
-    monitor_->Stop(__func__);
-
-    return nodes.front();
-  }
-
-  void UpdatePredictionCache(DMatrix const *data, linalg::VectorView<float> out_preds) const {
-    monitor_->Start(__func__);
-    // Caching prediction seems redundant for approx tree method, as sketching takes up
-    // majority of training time.
-    CHECK_EQ(out_preds.Size(), data->Info().num_row_);
-    UpdatePredictionCacheImpl(ctx_, p_last_tree_, partitioner_, evaluator_, param_, out_preds);
-    monitor_->Stop(__func__);
-  }
-
-  void BuildHistogram(DMatrix *p_fmat, RegTree *p_tree,
-                      std::vector<CPUExpandEntry> const &valid_candidates,
-                      std::vector<GradientPair> const &gpair, common::Span<float> hess) {
-    monitor_->Start(__func__);
-    std::vector<CPUExpandEntry> nodes_to_build;
-    std::vector<CPUExpandEntry> nodes_to_sub;
-
-    for (auto const &c : valid_candidates) {
-      auto left_nidx = (*p_tree)[c.nid].LeftChild();
-      auto right_nidx = (*p_tree)[c.nid].RightChild();
-      auto fewer_right = c.split.right_sum.GetHess() < c.split.left_sum.GetHess();
-
-      auto build_nidx = left_nidx;
-      auto subtract_nidx = right_nidx;
-      if (fewer_right) {
-        std::swap(build_nidx, subtract_nidx);
-      }
-      nodes_to_build.push_back(CPUExpandEntry{build_nidx, p_tree->GetDepth(build_nidx), {}});
-      nodes_to_sub.push_back(CPUExpandEntry{subtract_nidx, p_tree->GetDepth(subtract_nidx), {}});
-    }
-
-    size_t i = 0;
-    auto space = ConstructHistSpace(partitioner_, nodes_to_build);
-    for (auto const &page : p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess))) {
-      histogram_builder_.BuildHist(i, space, page, p_tree, partitioner_.at(i).Partitions(),
-                                   nodes_to_build, nodes_to_sub, gpair);
-      i++;
-    }
-    monitor_->Stop(__func__);
-  }
-
- public:
-  explicit GloablApproxBuilder(TrainParam param, MetaInfo const &info, GenericParameter const *ctx,
-                               std::shared_ptr<common::ColumnSampler> column_sampler, ObjInfo task,
-                               common::Monitor *monitor)
-      : param_{std::move(param)},
-        col_sampler_{std::move(column_sampler)},
-        evaluator_{param_, info, ctx->Threads(), col_sampler_, task},
-        ctx_{ctx},
-        monitor_{monitor} {}
-
-  void UpdateTree(DMatrix *p_fmat, std::vector<GradientPair> const &gpair, common::Span<float> hess,
-                  RegTree *p_tree, std::vector<RowIndexCache> *p_out_row_indices) {
-    p_last_tree_ = p_tree;
-    this->InitData(p_fmat, hess);
-
-    Driver<CPUExpandEntry> driver(static_cast<TrainParam::TreeGrowPolicy>(param_.grow_policy));
-    auto &tree = *p_tree;
-    driver.Push({this->InitRoot(p_fmat, gpair, hess, p_tree)});
-    bst_node_t num_leaves{1};
-    auto expand_set = driver.Pop();
-
-    /**
-     * Note for update position
-     * Root:
-     *   Not applied: No need to update position as initialization has got all the rows ordered.
-     *   Applied: Update position is run on applied nodes so the rows are partitioned.
-     * Non-root:
-     *   Not applied: That node is root of the subtree, same rule as root.
-     *   Applied: Ditto
-     */
-
-    while (!expand_set.empty()) {
-      // candidates that can be further splited.
-      std::vector<CPUExpandEntry> valid_candidates;
-      // candidates that can be applied.
-      std::vector<CPUExpandEntry> applied;
-      for (auto const &candidate : expand_set) {
-        if (!candidate.IsValid(param_, num_leaves)) {
-          continue;
-        }
-        evaluator_.ApplyTreeSplit(candidate, p_tree);
-        applied.push_back(candidate);
-        num_leaves++;
-        int left_child_nidx = tree[candidate.nid].LeftChild();
-        if (CPUExpandEntry::ChildIsValid(param_, p_tree->GetDepth(left_child_nidx), num_leaves)) {
-          valid_candidates.emplace_back(candidate);
-        }
-      }
-
-      monitor_->Start("UpdatePosition");
-      size_t page_id = 0;
-      for (auto const &page : p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess))) {
-        partitioner_.at(page_id).UpdatePosition(ctx_, page, applied, p_tree);
-        page_id++;
-      }
-      monitor_->Stop("UpdatePosition");
-
-      std::vector<CPUExpandEntry> best_splits;
-      if (!valid_candidates.empty()) {
-        this->BuildHistogram(p_fmat, p_tree, valid_candidates, gpair, hess);
-        for (auto const &candidate : valid_candidates) {
-          int left_child_nidx = tree[candidate.nid].LeftChild();
-          int right_child_nidx = tree[candidate.nid].RightChild();
-          CPUExpandEntry l_best{left_child_nidx, tree.GetDepth(left_child_nidx), {}};
-          CPUExpandEntry r_best{right_child_nidx, tree.GetDepth(right_child_nidx), {}};
-          best_splits.push_back(l_best);
-          best_splits.push_back(r_best);
-        }
-        auto const &histograms = histogram_builder_.Histogram();
-        auto ft = p_fmat->Info().feature_types.ConstHostSpan();
-        monitor_->Start("EvaluateSplits");
-        evaluator_.EvaluateSplits(histograms, feature_values_, ft, *p_tree, &best_splits);
-        monitor_->Stop("EvaluateSplits");
-      }
-      driver.Push(best_splits.begin(), best_splits.end());
-      expand_set = driver.Pop();
-    }
-
-    CHECK(p_out_row_indices->empty());
-    for (auto const &part : partitioner_) {
-      p_out_row_indices->emplace_back();
-      auto row_set = part.Partitions();
-      auto n_leaf = row_set.Size();
-      // fixme: subsample
-      auto &h_row_index = p_out_row_indices->back().row_index.HostVector();
-
-      auto begin = row_set.Data()->data();
-      for (auto node : row_set) {
-        CHECK(node.begin);
-        CHECK(tree[node.node_id].IsLeaf());
-        size_t offset = node.begin - begin;
-        auto size = node.Size();
-        auto seg = RowIndexCache::Segment{offset, size, node.node_id};
-        size_t k = seg.begin;
-        for (auto idx = node.begin; idx != node.end; ++idx) {
-          h_row_index[k] = *idx;
-        }
-        p_out_row_indices->back().indptr.push_back(seg);
-      }
-    }
-  }
-};
-
-/**
- * \brief Implementation for the approx tree method.  It constructs quantile for every
- *        iteration.
- */
-class GlobalApproxUpdater : public TreeUpdater {
-  TrainParam param_;
-  common::Monitor monitor_;
-  CPUHistMakerTrainParam hist_param_;
-  // specializations for different histogram precision.
-  std::unique_ptr<GloablApproxBuilder<float>> f32_impl_;
-  std::unique_ptr<GloablApproxBuilder<double>> f64_impl_;
-  // pointer to the last DMatrix, used for update prediction cache.
-  DMatrix *cached_{nullptr};
-  std::shared_ptr<common::ColumnSampler> column_sampler_ =
-      std::make_shared<common::ColumnSampler>();
-  // cache for row partitions
-  std::vector<std::vector<RowIndexCache>> row_set_collection_;
-  ObjInfo task_;
-
- public:
-  explicit GlobalApproxUpdater(ObjInfo task) : task_{task} { monitor_.Init(__func__); }
-
-  void Configure(const Args &args) override {
-    param_.UpdateAllowUnknown(args);
-    hist_param_.UpdateAllowUnknown(args);
-  }
-  void LoadConfig(Json const &in) override {
-    auto const &config = get<Object const>(in);
-    FromJson(config.at("train_param"), &this->param_);
-    FromJson(config.at("hist_param"), &this->hist_param_);
-  }
-  void SaveConfig(Json *p_out) const override {
-    auto &out = *p_out;
-    out["train_param"] = ToJson(param_);
-    out["hist_param"] = ToJson(hist_param_);
-  }
-
-  void InitData(TrainParam const &param, HostDeviceVector<GradientPair> const *gpair,
-                std::vector<GradientPair> *sampled) {
-    auto const &h_gpair = gpair->ConstHostVector();
-    sampled->resize(h_gpair.size());
-    std::copy(h_gpair.cbegin(), h_gpair.cend(), sampled->begin());
-    auto &rnd = common::GlobalRandom();
-    row_set_collection_.clear();
-
-    if (param.subsample != 1.0) {
-      CHECK(param.sampling_method != TrainParam::kGradientBased)
-          << "Gradient based sampling is not supported for approx tree method.";
-      std::bernoulli_distribution coin_flip(param.subsample);
-      std::transform(sampled->begin(), sampled->end(), sampled->begin(), [&](GradientPair &g) {
-        if (coin_flip(rnd)) {
-          return g;
-        } else {
-          return GradientPair{};
-        }
-      });
-    }
-  }
-
-  char const *Name() const override { return "grow_histmaker"; }
-
-  void Update(HostDeviceVector<GradientPair> *gpair, DMatrix *m,
-              const std::vector<RegTree *> &trees) override {
-    float lr = param_.learning_rate;
-    param_.learning_rate = lr / trees.size();
-
-    if (hist_param_.single_precision_histogram) {
-      f32_impl_ = std::make_unique<GloablApproxBuilder<float>>(param_, m->Info(), ctx_,
-                                                               column_sampler_, task_, &monitor_);
-    } else {
-      f64_impl_ = std::make_unique<GloablApproxBuilder<double>>(param_, m->Info(), ctx_,
-                                                                column_sampler_, task_, &monitor_);
-    }
-
-    std::vector<GradientPair> h_gpair;
-    InitData(param_, gpair, &h_gpair);
-    // Obtain the hessian values for weighted sketching
-    std::vector<float> hess(h_gpair.size());
-    std::transform(h_gpair.begin(), h_gpair.end(), hess.begin(),
-                   [](auto g) { return g.GetHess(); });
-
-    cached_ = m;
-
-    for (auto p_tree : trees) {
-      row_set_collection_.emplace_back();
-      auto &row_indices = row_set_collection_.back();
-      if (hist_param_.single_precision_histogram) {
-        this->f32_impl_->UpdateTree(m, h_gpair, hess, p_tree, &row_indices);
-      } else {
-        this->f64_impl_->UpdateTree(m, h_gpair, hess, p_tree, &row_indices);
-      }
-    }
-    param_.learning_rate = lr;
-  }
-
-  bool UpdatePredictionCache(const DMatrix *data, linalg::VectorView<float> out_preds) override {
-    if (data != cached_ || (!this->f32_impl_ && !this->f64_impl_)) {
-      return false;
-    }
-
-    if (hist_param_.single_precision_histogram) {
-      this->f32_impl_->UpdatePredictionCache(data, out_preds);
-    } else {
-      this->f64_impl_->UpdatePredictionCache(data, out_preds);
-    }
-    return true;
-  }
-
-  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
-    return row_set_collection_.at(tree_idx);
-  }
-};
-
-DMLC_REGISTRY_FILE_TAG(grow_histmaker);
-
-XGBOOST_REGISTER_TREE_UPDATER(GlobalHistMaker, "grow_histmaker")
-    .describe(
-        "Tree constructor that uses approximate histogram construction "
-        "for each node.")
-    .set_body([](ObjInfo task) { return new GlobalApproxUpdater(task); });
-}  // namespace tree
-}  // namespace xgboost
+DoBoost(DMatrix *p_fmat, HostDeviceVector<GradientPair> *in_gpair, PredictionCacheEntry *predt, const ObjFunction *obj) -> void
+\brief Carry out one iteration of boosting
diff --git a/tests/cpp/common/test_stats.cc b/tests/cpp/common/test_stats.cc
new file mode 100644
index 000000000000..94bf53e7288a
--- /dev/null
+++ b/tests/cpp/common/test_stats.cc
@@ -0,0 +1,14 @@
+#include <gtest/gtest.h>
+#include <xgboost/generic_parameters.h>
+
+#include "../../../src/common/stats.h"
+
+namespace xgboost {
+namespace common {
+TEST(Stats, Percentil) {
+  linalg::Tensor<float, 1> arr({21, 15, 50, 40, 35}, {5}, Context::kCpuId);
+  auto percentile = Percentile(40.f, arr.HostView());
+  std::cout << percentile << std::endl;
+}
+}  // namespace common
+}  // namespace xgboost
diff --git a/tests/cpp/gbm/test_gbtree.cc b/tests/cpp/gbm/test_gbtree.cc
index c416d134307c..d703c5dbac4a 100644
--- a/tests/cpp/gbm/test_gbtree.cc
+++ b/tests/cpp/gbm/test_gbtree.cc
@@ -69,13 +69,13 @@ TEST(GBTree, PredictionCache) {
   auto p_m = RandomDataGenerator{kRows, kCols, 0}.GenerateDMatrix();
   auto gpair = GenerateRandomGradients(kRows);
   PredictionCacheEntry out_predictions;
-  gbtree.DoBoost(p_m.get(), &gpair, &out_predictions);
+  gbtree.DoBoost(p_m.get(), &gpair, &out_predictions, nullptr);
 
   gbtree.PredictBatch(p_m.get(), &out_predictions, false, 0, 0);
   ASSERT_EQ(1, out_predictions.version);
   std::vector<float> first_iter = out_predictions.predictions.HostVector();
   // Add 1 more boosted round
-  gbtree.DoBoost(p_m.get(), &gpair, &out_predictions);
+  gbtree.DoBoost(p_m.get(), &gpair, &out_predictions, nullptr);
   gbtree.PredictBatch(p_m.get(), &out_predictions, false, 0, 0);
   ASSERT_EQ(2, out_predictions.version);
   // Update the cache for all rounds
@@ -83,7 +83,7 @@ TEST(GBTree, PredictionCache) {
   gbtree.PredictBatch(p_m.get(), &out_predictions, false, 0, 0);
   ASSERT_EQ(2, out_predictions.version);
 
-  gbtree.DoBoost(p_m.get(), &gpair, &out_predictions);
+  gbtree.DoBoost(p_m.get(), &gpair, &out_predictions, nullptr);
   // drop the cache.
   gbtree.PredictBatch(p_m.get(), &out_predictions, false, 1, 2);
   ASSERT_EQ(0, out_predictions.version);
diff --git a/tests/cpp/helpers.cc b/tests/cpp/helpers.cc
index 05c138781e0d..68faa09642ed 100644
--- a/tests/cpp/helpers.cc
+++ b/tests/cpp/helpers.cc
@@ -548,7 +548,7 @@ std::unique_ptr<GradientBooster> CreateTrainedGBM(
 
   PredictionCacheEntry predts;
 
-  gbm->DoBoost(p_dmat.get(), &gpair, &predts);
+  gbm->DoBoost(p_dmat.get(), &gpair, &predts, nullptr);
 
   return gbm;
 }
diff --git a/tests/cpp/predictor/test_cpu_predictor.cc b/tests/cpp/predictor/test_cpu_predictor.cc
index 1a466ed3ff10..279aacea54b7 100644
--- a/tests/cpp/predictor/test_cpu_predictor.cc
+++ b/tests/cpp/predictor/test_cpu_predictor.cc
@@ -222,7 +222,7 @@ void TestUpdatePredictionCache(bool use_subsampling) {
   PredictionCacheEntry predtion_cache;
   predtion_cache.predictions.Resize(kRows*kClasses, 0);
   // after one training iteration predtion_cache is filled with cached in QuantileHistMaker::Builder prediction values
-  gbm->DoBoost(dmat.get(), &gpair, &predtion_cache);
+  gbm->DoBoost(dmat.get(), &gpair, &predtion_cache, nullptr);
 
   PredictionCacheEntry out_predictions;
   // perform fair prediction on the same input data, should be equal to cached result

From ef5a14194589e62f847f3960dcf3df4eda844ddb Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 15:25:38 +0800
Subject: [PATCH 007/124] Fixes.

---
 src/common/stats.h              |  26 ++-
 src/gbm/gbtree.cc               |  16 +-
 src/objective/regression_obj.cc |  11 +-
 src/tree/updater_approx.cc      | 383 +++++++++++++++++++++++++++++++-
 tests/cpp/common/test_stats.cc  |   5 +-
 5 files changed, 416 insertions(+), 25 deletions(-)

diff --git a/src/common/stats.h b/src/common/stats.h
index c6fae6bb0a7c..b2cea3fc5aac 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -1,3 +1,8 @@
+/*!
+ * Copyright 2022 by XGBoost Contributors
+ */
+#ifndef XGBOOST_COMMON_STATS_H_
+#define XGBOOST_COMMON_STATS_H_
 #include <limits>
 #include <vector>
 
@@ -6,26 +11,33 @@
 
 namespace xgboost {
 namespace common {
-float Percentile(float percentile, linalg::TensorView<float const, 1> arr) {
-  size_t n = arr.Shape(0);
+inline float Percentile(float percentile, Span<size_t const> index,
+                        linalg::TensorView<float const, 1> arr) {
+  size_t n = index.size();
   if (n == 0) {
     return std::numeric_limits<float>::quiet_NaN();
   }
-  std::vector<size_t> sorted_idx{ArgSort(arr)};
+  std::vector<size_t> sorted_idx(index.size());
+  std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
+  std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
+                   [&](size_t l, size_t r) { return arr(index(l)) < arr(index(r)); });
+
+  auto val = [&](size_t i) { return arr(index(sorted_idx[i])); };
 
   if (percentile <= (1 / (n + 1))) {
-    return arr(sorted_idx.front());
+    return val(0);
   }
   if (percentile >= (n / (n + 1))) {
-    return arr(sorted_idx.back());
+    return val(sorted_idx.size() - 1);
   }
   double x = percentile * static_cast<double>((n + 1));
   double k = std::floor(x);
   double d = x - k;
 
-  auto v0 = arr(sorted_idx[static_cast<size_t>(k)]);
-  auto v1 = arr(sorted_idx[static_cast<size_t>(k) + 1]);
+  auto v0 = val(static_cast<size_t>(k));
+  auto v1 = val(static_cast<size_t>(k) + 1);
   return v0 + d * (v1 - v0);
 }
 }  // namespace common
 }  // namespace xgboost
+#endif  // XGBOOST_COMMON_STATS_H_
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index 372559d41205..8b801b7a6729 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -271,12 +271,16 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
     }
   }
 
-  bst_group_t gidx {0};
-  for (auto& tree_group : new_trees) {
-    for (size_t t = 0; t < tree_group.size(); ++t) {
-      auto row_idx = updaters_.back()->GetRowIndexCache(t);
-      auto target = p_fmat->Info().labels.Shape(1) > 1 ? gidx : 0;
-      obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), gidx, tree_group[t].get());
+  if (obj) {
+    // Update tree leaf value at the end of boosting.
+    bst_group_t gidx{0};
+    auto targets = obj->Targets(p_fmat->Info());
+    for (auto& tree_group : new_trees) {
+      for (size_t t = 0; t < tree_group.size(); ++t) {
+        auto row_idx = updaters_.back()->GetRowIndexCache(t);
+        auto target = targets > 1 ? gidx : 0;
+        obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), gidx, tree_group[t].get());
+      }
     }
   }
   monitor_.Stop("BoostNewTrees");
diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index 75afd78963c4..8f3ad23a27e0 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -7,6 +7,7 @@
 #include <dmlc/registry.h>
 
 #include "../common/linalg_op.h"
+#include "../common/stats.h"
 #include "rabit/rabit.h"
 #include "xgboost/data.h"
 #include "xgboost/objective.h"
@@ -48,13 +49,6 @@ float WeightedQuantile(float quantile, common::Span<size_t const> row_set,
   }
 };
 
-float Quantile(float quantile, common::Span<size_t const> row_set, linalg::VectorView<float const> labels) {
-  float result;
-  LOG(FATAL) << "Not implemented";
-  // fixme: pick an algorithm from R quantile.
-  return result;
-}
-
 void UpdateTreeLeafHost(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
                         uint32_t target, float alpha, RegTree* p_tree) {
   auto& tree = *p_tree;
@@ -65,7 +59,8 @@ void UpdateTreeLeafHost(common::Span<RowIndexCache const> row_index, MetaInfo co
       auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
       float q{0};
       if (info.weights_.Empty()) {
-        q = Quantile(alpha, h_row_set, info.labels.HostView().Slice(linalg::All(), target));
+        q = common::Percentile(alpha, h_row_set,
+                               info.labels.HostView().Slice(linalg::All(), target));
       } else {
         q = WeightedQuantile(alpha, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
                              linalg::MakeVec(&info.weights_));
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index fd01034815bc..d3f98baf0388 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -1,2 +1,381 @@
-DoBoost(DMatrix *p_fmat, HostDeviceVector<GradientPair> *in_gpair, PredictionCacheEntry *predt, const ObjFunction *obj) -> void
-\brief Carry out one iteration of boosting
+/*!
+ * Copyright 2021-2022 XGBoost contributors
+ *
+ * \brief Implementation for the approx tree method.
+ */
+#include "updater_approx.h"
+
+#include <algorithm>
+#include <memory>
+#include <vector>
+
+#include "../common/random.h"
+#include "../data/gradient_index.h"
+#include "constraints.h"
+#include "driver.h"
+#include "hist/evaluate_splits.h"
+#include "hist/histogram.h"
+#include "hist/param.h"
+#include "param.h"
+#include "xgboost/base.h"
+#include "xgboost/json.h"
+#include "xgboost/tree_updater.h"
+
+namespace xgboost {
+namespace tree {
+
+DMLC_REGISTRY_FILE_TAG(updater_approx);
+
+namespace {
+// Return the BatchParam used by DMatrix.
+template <typename GradientSumT>
+auto BatchSpec(TrainParam const &p, common::Span<float> hess,
+               HistEvaluator<GradientSumT, CPUExpandEntry> const &evaluator) {
+  return BatchParam{p.max_bin, hess, !evaluator.Task().const_hess};
+}
+
+auto BatchSpec(TrainParam const &p, common::Span<float> hess) {
+  return BatchParam{p.max_bin, hess, false};
+}
+}  // anonymous namespace
+
+template <typename GradientSumT>
+class GloablApproxBuilder {
+ protected:
+  TrainParam param_;
+  std::shared_ptr<common::ColumnSampler> col_sampler_;
+  HistEvaluator<GradientSumT, CPUExpandEntry> evaluator_;
+  HistogramBuilder<GradientSumT, CPUExpandEntry> histogram_builder_;
+  GenericParameter const *ctx_;
+
+  std::vector<ApproxRowPartitioner> partitioner_;
+  // Pointer to last updated tree, used for update prediction cache.
+  RegTree *p_last_tree_{nullptr};
+  common::Monitor *monitor_;
+  size_t n_batches_{0};
+  // Cache for histogram cuts.
+  common::HistogramCuts feature_values_;
+
+ public:
+  void InitData(DMatrix *p_fmat, common::Span<float> hess) {
+    monitor_->Start(__func__);
+
+    n_batches_ = 0;
+    int32_t n_total_bins = 0;
+    partitioner_.clear();
+    // Generating the GHistIndexMatrix is quite slow, is there a way to speed it up?
+    for (auto const &page :
+         p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess, evaluator_))) {
+      if (n_total_bins == 0) {
+        n_total_bins = page.cut.TotalBins();
+        feature_values_ = page.cut;
+      } else {
+        CHECK_EQ(n_total_bins, page.cut.TotalBins());
+      }
+      partitioner_.emplace_back(page.Size(), page.base_rowid);
+      n_batches_++;
+    }
+
+    histogram_builder_.Reset(n_total_bins, BatchSpec(param_, hess), ctx_->Threads(), n_batches_,
+                             rabit::IsDistributed());
+    monitor_->Stop(__func__);
+  }
+
+  CPUExpandEntry InitRoot(DMatrix *p_fmat, std::vector<GradientPair> const &gpair,
+                          common::Span<float> hess, RegTree *p_tree) {
+    monitor_->Start(__func__);
+    CPUExpandEntry best;
+    best.nid = RegTree::kRoot;
+    best.depth = 0;
+    GradStats root_sum;
+    for (auto const &g : gpair) {
+      root_sum.Add(g);
+    }
+    rabit::Allreduce<rabit::op::Sum, double>(reinterpret_cast<double *>(&root_sum), 2);
+    std::vector<CPUExpandEntry> nodes{best};
+    size_t i = 0;
+    auto space = ConstructHistSpace(partitioner_, nodes);
+    for (auto const &page : p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess))) {
+      histogram_builder_.BuildHist(i, space, page, p_tree, partitioner_.at(i).Partitions(), nodes,
+                                   {}, gpair);
+      i++;
+    }
+
+    auto weight = evaluator_.InitRoot(root_sum);
+    p_tree->Stat(RegTree::kRoot).sum_hess = root_sum.GetHess();
+    p_tree->Stat(RegTree::kRoot).base_weight = weight;
+    (*p_tree)[RegTree::kRoot].SetLeaf(param_.learning_rate * weight);
+
+    auto const &histograms = histogram_builder_.Histogram();
+    auto ft = p_fmat->Info().feature_types.ConstHostSpan();
+    evaluator_.EvaluateSplits(histograms, feature_values_, ft, *p_tree, &nodes);
+    monitor_->Stop(__func__);
+
+    return nodes.front();
+  }
+
+  void UpdatePredictionCache(DMatrix const *data, linalg::VectorView<float> out_preds) const {
+    monitor_->Start(__func__);
+    // Caching prediction seems redundant for approx tree method, as sketching takes up
+    // majority of training time.
+    CHECK_EQ(out_preds.Size(), data->Info().num_row_);
+    UpdatePredictionCacheImpl(ctx_, p_last_tree_, partitioner_, evaluator_, param_, out_preds);
+    monitor_->Stop(__func__);
+  }
+
+  void BuildHistogram(DMatrix *p_fmat, RegTree *p_tree,
+                      std::vector<CPUExpandEntry> const &valid_candidates,
+                      std::vector<GradientPair> const &gpair, common::Span<float> hess) {
+    monitor_->Start(__func__);
+    std::vector<CPUExpandEntry> nodes_to_build;
+    std::vector<CPUExpandEntry> nodes_to_sub;
+
+    for (auto const &c : valid_candidates) {
+      auto left_nidx = (*p_tree)[c.nid].LeftChild();
+      auto right_nidx = (*p_tree)[c.nid].RightChild();
+      auto fewer_right = c.split.right_sum.GetHess() < c.split.left_sum.GetHess();
+
+      auto build_nidx = left_nidx;
+      auto subtract_nidx = right_nidx;
+      if (fewer_right) {
+        std::swap(build_nidx, subtract_nidx);
+      }
+      nodes_to_build.push_back(CPUExpandEntry{build_nidx, p_tree->GetDepth(build_nidx), {}});
+      nodes_to_sub.push_back(CPUExpandEntry{subtract_nidx, p_tree->GetDepth(subtract_nidx), {}});
+    }
+
+    size_t i = 0;
+    auto space = ConstructHistSpace(partitioner_, nodes_to_build);
+    for (auto const &page : p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess))) {
+      histogram_builder_.BuildHist(i, space, page, p_tree, partitioner_.at(i).Partitions(),
+                                   nodes_to_build, nodes_to_sub, gpair);
+      i++;
+    }
+    monitor_->Stop(__func__);
+  }
+
+ public:
+  explicit GloablApproxBuilder(TrainParam param, MetaInfo const &info, GenericParameter const *ctx,
+                               std::shared_ptr<common::ColumnSampler> column_sampler, ObjInfo task,
+                               common::Monitor *monitor)
+      : param_{std::move(param)},
+        col_sampler_{std::move(column_sampler)},
+        evaluator_{param_, info, ctx->Threads(), col_sampler_, task},
+        ctx_{ctx},
+        monitor_{monitor} {}
+
+  void UpdateTree(DMatrix *p_fmat, std::vector<GradientPair> const &gpair, common::Span<float> hess,
+                  RegTree *p_tree, std::vector<RowIndexCache> *p_out_row_indices) {
+    p_last_tree_ = p_tree;
+    this->InitData(p_fmat, hess);
+
+    Driver<CPUExpandEntry> driver(static_cast<TrainParam::TreeGrowPolicy>(param_.grow_policy));
+    auto &tree = *p_tree;
+    driver.Push({this->InitRoot(p_fmat, gpair, hess, p_tree)});
+    bst_node_t num_leaves{1};
+    auto expand_set = driver.Pop();
+
+    /**
+     * Note for update position
+     * Root:
+     *   Not applied: No need to update position as initialization has got all the rows ordered.
+     *   Applied: Update position is run on applied nodes so the rows are partitioned.
+     * Non-root:
+     *   Not applied: That node is root of the subtree, same rule as root.
+     *   Applied: Ditto
+     */
+
+    while (!expand_set.empty()) {
+      // candidates that can be further splited.
+      std::vector<CPUExpandEntry> valid_candidates;
+      // candidates that can be applied.
+      std::vector<CPUExpandEntry> applied;
+      for (auto const &candidate : expand_set) {
+        if (!candidate.IsValid(param_, num_leaves)) {
+          continue;
+        }
+        evaluator_.ApplyTreeSplit(candidate, p_tree);
+        applied.push_back(candidate);
+        num_leaves++;
+        int left_child_nidx = tree[candidate.nid].LeftChild();
+        if (CPUExpandEntry::ChildIsValid(param_, p_tree->GetDepth(left_child_nidx), num_leaves)) {
+          valid_candidates.emplace_back(candidate);
+        }
+      }
+
+      monitor_->Start("UpdatePosition");
+      size_t page_id = 0;
+      for (auto const &page : p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess))) {
+        partitioner_.at(page_id).UpdatePosition(ctx_, page, applied, p_tree);
+        page_id++;
+      }
+      monitor_->Stop("UpdatePosition");
+
+      std::vector<CPUExpandEntry> best_splits;
+      if (!valid_candidates.empty()) {
+        this->BuildHistogram(p_fmat, p_tree, valid_candidates, gpair, hess);
+        for (auto const &candidate : valid_candidates) {
+          int left_child_nidx = tree[candidate.nid].LeftChild();
+          int right_child_nidx = tree[candidate.nid].RightChild();
+          CPUExpandEntry l_best{left_child_nidx, tree.GetDepth(left_child_nidx), {}};
+          CPUExpandEntry r_best{right_child_nidx, tree.GetDepth(right_child_nidx), {}};
+          best_splits.push_back(l_best);
+          best_splits.push_back(r_best);
+        }
+        auto const &histograms = histogram_builder_.Histogram();
+        auto ft = p_fmat->Info().feature_types.ConstHostSpan();
+        monitor_->Start("EvaluateSplits");
+        evaluator_.EvaluateSplits(histograms, feature_values_, ft, *p_tree, &best_splits);
+        monitor_->Stop("EvaluateSplits");
+      }
+      driver.Push(best_splits.begin(), best_splits.end());
+      expand_set = driver.Pop();
+    }
+
+    CHECK(p_out_row_indices->empty());
+    for (auto const &part : partitioner_) {
+      auto row_set = part.Partitions();
+      auto n_leaf = row_set.Size();
+      p_out_row_indices->emplace_back(ctx_, n_leaf, row_set.Size());
+      // fixme: subsample
+      auto &h_row_index = p_out_row_indices->back().row_index.HostVector();
+
+      auto begin = row_set.Data()->data();
+      for (auto node : row_set) {
+        CHECK(node.begin);
+        CHECK(tree[node.node_id].IsLeaf());
+        size_t offset = node.begin - begin;
+        auto size = node.Size();
+        auto seg = RowIndexCache::Segment{offset, size, node.node_id};
+        size_t k = seg.begin;
+        for (auto idx = node.begin; idx != node.end; ++idx) {
+          h_row_index[k] = *idx;
+        }
+        p_out_row_indices->back().indptr.push_back(seg);
+      }
+    }
+  }
+};
+
+/**
+ * \brief Implementation for the approx tree method.  It constructs quantile for every
+ *        iteration.
+ */
+class GlobalApproxUpdater : public TreeUpdater {
+  TrainParam param_;
+  common::Monitor monitor_;
+  CPUHistMakerTrainParam hist_param_;
+  // specializations for different histogram precision.
+  std::unique_ptr<GloablApproxBuilder<float>> f32_impl_;
+  std::unique_ptr<GloablApproxBuilder<double>> f64_impl_;
+  // pointer to the last DMatrix, used for update prediction cache.
+  DMatrix *cached_{nullptr};
+  std::shared_ptr<common::ColumnSampler> column_sampler_ =
+      std::make_shared<common::ColumnSampler>();
+  // cache for row partitions
+  std::vector<std::vector<RowIndexCache>> row_set_collection_;
+  ObjInfo task_;
+
+ public:
+  explicit GlobalApproxUpdater(ObjInfo task) : task_{task} { monitor_.Init(__func__); }
+
+  void Configure(const Args &args) override {
+    param_.UpdateAllowUnknown(args);
+    hist_param_.UpdateAllowUnknown(args);
+  }
+  void LoadConfig(Json const &in) override {
+    auto const &config = get<Object const>(in);
+    FromJson(config.at("train_param"), &this->param_);
+    FromJson(config.at("hist_param"), &this->hist_param_);
+  }
+  void SaveConfig(Json *p_out) const override {
+    auto &out = *p_out;
+    out["train_param"] = ToJson(param_);
+    out["hist_param"] = ToJson(hist_param_);
+  }
+
+  void InitData(TrainParam const &param, HostDeviceVector<GradientPair> const *gpair,
+                std::vector<GradientPair> *sampled) {
+    auto const &h_gpair = gpair->ConstHostVector();
+    sampled->resize(h_gpair.size());
+    std::copy(h_gpair.cbegin(), h_gpair.cend(), sampled->begin());
+    auto &rnd = common::GlobalRandom();
+    row_set_collection_.clear();
+
+    if (param.subsample != 1.0) {
+      CHECK(param.sampling_method != TrainParam::kGradientBased)
+          << "Gradient based sampling is not supported for approx tree method.";
+      std::bernoulli_distribution coin_flip(param.subsample);
+      std::transform(sampled->begin(), sampled->end(), sampled->begin(), [&](GradientPair &g) {
+        if (coin_flip(rnd)) {
+          return g;
+        } else {
+          return GradientPair{};
+        }
+      });
+    }
+  }
+
+  char const *Name() const override { return "grow_histmaker"; }
+
+  void Update(HostDeviceVector<GradientPair> *gpair, DMatrix *m,
+              const std::vector<RegTree *> &trees) override {
+    float lr = param_.learning_rate;
+    param_.learning_rate = lr / trees.size();
+
+    if (hist_param_.single_precision_histogram) {
+      f32_impl_ = std::make_unique<GloablApproxBuilder<float>>(param_, m->Info(), ctx_,
+                                                               column_sampler_, task_, &monitor_);
+    } else {
+      f64_impl_ = std::make_unique<GloablApproxBuilder<double>>(param_, m->Info(), ctx_,
+                                                                column_sampler_, task_, &monitor_);
+    }
+
+    std::vector<GradientPair> h_gpair;
+    InitData(param_, gpair, &h_gpair);
+    // Obtain the hessian values for weighted sketching
+    std::vector<float> hess(h_gpair.size());
+    std::transform(h_gpair.begin(), h_gpair.end(), hess.begin(),
+                   [](auto g) { return g.GetHess(); });
+
+    cached_ = m;
+
+    for (auto p_tree : trees) {
+      row_set_collection_.emplace_back();
+      auto &row_indices = row_set_collection_.back();
+      if (hist_param_.single_precision_histogram) {
+        this->f32_impl_->UpdateTree(m, h_gpair, hess, p_tree, &row_indices);
+      } else {
+        this->f64_impl_->UpdateTree(m, h_gpair, hess, p_tree, &row_indices);
+      }
+    }
+    param_.learning_rate = lr;
+  }
+
+  bool UpdatePredictionCache(const DMatrix *data, linalg::VectorView<float> out_preds) override {
+    if (data != cached_ || (!this->f32_impl_ && !this->f64_impl_)) {
+      return false;
+    }
+
+    if (hist_param_.single_precision_histogram) {
+      this->f32_impl_->UpdatePredictionCache(data, out_preds);
+    } else {
+      this->f64_impl_->UpdatePredictionCache(data, out_preds);
+    }
+    return true;
+  }
+
+  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
+    return row_set_collection_.at(tree_idx);
+  }
+};
+
+DMLC_REGISTRY_FILE_TAG(grow_histmaker);
+
+XGBOOST_REGISTER_TREE_UPDATER(GlobalHistMaker, "grow_histmaker")
+    .describe(
+        "Tree constructor that uses approximate histogram construction "
+        "for each node.")
+    .set_body([](ObjInfo task) { return new GlobalApproxUpdater(task); });
+}  // namespace tree
+}  // namespace xgboost
diff --git a/tests/cpp/common/test_stats.cc b/tests/cpp/common/test_stats.cc
index 94bf53e7288a..e61a6e42bc27 100644
--- a/tests/cpp/common/test_stats.cc
+++ b/tests/cpp/common/test_stats.cc
@@ -6,8 +6,9 @@
 namespace xgboost {
 namespace common {
 TEST(Stats, Percentil) {
-  linalg::Tensor<float, 1> arr({21, 15, 50, 40, 35}, {5}, Context::kCpuId);
-  auto percentile = Percentile(40.f, arr.HostView());
+  linalg::Tensor<float, 1> arr({21, 0, 15, 50, 40, 0, 35}, {5}, Context::kCpuId);
+  std::vector<size_t> index{0, 2, 3, 4, 6};
+  auto percentile = Percentile(40.f, Span<size_t const>{index}, arr.HostView());
   std::cout << percentile << std::endl;
 }
 }  // namespace common

From d19e66748df4e23720002e4c3da0267c2e42fb1b Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 16:02:11 +0800
Subject: [PATCH 008/124] Test.

---
 src/common/stats.h             | 30 ++++++++++++++++++++++--------
 src/gbm/gbtree.cc              |  8 +++-----
 tests/cpp/common/test_stats.cc | 12 +++++++++---
 3 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/src/common/stats.h b/src/common/stats.h
index b2cea3fc5aac..884c51324f22 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -11,12 +11,24 @@
 
 namespace xgboost {
 namespace common {
-inline float Percentile(float percentile, Span<size_t const> index,
+/**
+ * \brief Percentile with masked array using linear interpolation.
+ *
+ *   https://www.itl.nist.gov/div898/handbook/prc/section2/prc262.htm
+ *
+ * \param alpha percentile, must be in range [0, 1].
+ * \param index The index of valid elements in arr.
+ * \param arr   Input values.
+ *
+ * \return The result of interpolation.
+ */
+inline float Percentile(double alpha, Span<size_t const> index,
                         linalg::TensorView<float const, 1> arr) {
-  size_t n = index.size();
-  if (n == 0) {
+  CHECK(alpha >= 0 && alpha <= 1);
+  if (index.size() == 0) {
     return std::numeric_limits<float>::quiet_NaN();
   }
+  auto n = static_cast<double>(index.size());
   std::vector<size_t> sorted_idx(index.size());
   std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
   std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
@@ -24,15 +36,17 @@ inline float Percentile(float percentile, Span<size_t const> index,
 
   auto val = [&](size_t i) { return arr(index(sorted_idx[i])); };
 
-  if (percentile <= (1 / (n + 1))) {
+  if (alpha <= (1 / (n + 1))) {
     return val(0);
   }
-  if (percentile >= (n / (n + 1))) {
+  if (alpha >= (n / (n + 1))) {
     return val(sorted_idx.size() - 1);
   }
-  double x = percentile * static_cast<double>((n + 1));
-  double k = std::floor(x);
-  double d = x - k;
+
+  double x = alpha * static_cast<double>((n + 1));
+  double k = std::floor(x) - 1;
+  CHECK_GE(k, 0);
+  double d = (x - 1) - k;
 
   auto v0 = val(static_cast<size_t>(k));
   auto v1 = val(static_cast<size_t>(k) + 1);
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index 8b801b7a6729..74e13a342678 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -226,12 +226,9 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
   // Weird case that tree method is cpu-based but gpu_id is set.  Ideally we should let
   // `gpu_id` be the single source of determining what algorithms to run, but that will
   // break a lots of existing code.
-  auto device = tparam_.tree_method != TreeMethod::kGPUHist
-                    ? GenericParameter::kCpuId
-                    : ctx_->gpu_id;
+  auto device = tparam_.tree_method != TreeMethod::kGPUHist ? Context::kCpuId : ctx_->gpu_id;
   auto out = linalg::TensorView<float, 2>{
-      device == GenericParameter::kCpuId ? predt->predictions.HostSpan()
-                                         : predt->predictions.DeviceSpan(),
+      device == Context::kCpuId ? predt->predictions.HostSpan() : predt->predictions.DeviceSpan(),
       {static_cast<size_t>(p_fmat->Info().num_row_), static_cast<size_t>(ngroup)},
       device};
   CHECK_NE(ngroup, 0);
@@ -277,6 +274,7 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
     auto targets = obj->Targets(p_fmat->Info());
     for (auto& tree_group : new_trees) {
       for (size_t t = 0; t < tree_group.size(); ++t) {
+        // within a forest
         auto row_idx = updaters_.back()->GetRowIndexCache(t);
         auto target = targets > 1 ? gidx : 0;
         obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), gidx, tree_group[t].get());
diff --git a/tests/cpp/common/test_stats.cc b/tests/cpp/common/test_stats.cc
index e61a6e42bc27..49ab0fb98911 100644
--- a/tests/cpp/common/test_stats.cc
+++ b/tests/cpp/common/test_stats.cc
@@ -6,10 +6,16 @@
 namespace xgboost {
 namespace common {
 TEST(Stats, Percentil) {
-  linalg::Tensor<float, 1> arr({21, 0, 15, 50, 40, 0, 35}, {5}, Context::kCpuId);
+  linalg::Tensor<float, 1> arr({20.f, 0.f, 15.f, 50.f, 40.f, 0.f, 35.f}, {7}, Context::kCpuId);
   std::vector<size_t> index{0, 2, 3, 4, 6};
-  auto percentile = Percentile(40.f, Span<size_t const>{index}, arr.HostView());
-  std::cout << percentile << std::endl;
+  auto percentile = Percentile(0.40f, Span<size_t const>{index}, arr.HostView());
+  ASSERT_EQ(percentile, 26.0);
+
+  percentile = Percentile(0.20f, Span<size_t const>{index}, arr.HostView());
+  ASSERT_EQ(percentile, 16.0);
+
+  percentile = Percentile(0.10f, Span<size_t const>{index}, arr.HostView());
+  ASSERT_EQ(percentile, 15.0);
 }
 }  // namespace common
 }  // namespace xgboost

From 823648c208e89ec1cda48411669671a7a3408081 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 16:17:22 +0800
Subject: [PATCH 009/124] Use in boosting.

---
 include/xgboost/objective.h |  9 +++++++++
 src/gbm/gbtree.cc           | 33 ++++++++++++++++++---------------
 src/gbm/gbtree.h            |  3 +++
 3 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 385889d4784b..30ad4a2b0f71 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -91,6 +91,15 @@ class ObjFunction : public Configurable {
     return 1;
   }
 
+  /**
+   * \brief Update the leaf values after a tree is built. Needed for objectives with 0
+   *        hessian.
+   *
+   * \param row_index The index of rows for each output leaf.
+   * \param info MetaInfo providing labels and weights.
+   * \param target The index for target if we are training multi-target models, 0 otherwise.
+   * \param p_tree Tree that needs to be updated.
+   */
   virtual void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
                               uint32_t target, RegTree* p_tree) const {}
 
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index 74e13a342678..f661e1bcbda4 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -217,6 +217,21 @@ void CopyGradient(HostDeviceVector<GradientPair> const* in_gpair, int32_t n_thre
   }
 }
 
+void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, ObjFunction const* obj, size_t gidx,
+                            std::vector<std::unique_ptr<RegTree>>* p_trees) {
+  if (!obj) {
+    return;
+  }
+  auto& trees = *p_trees;
+  auto targets = obj->Targets(p_fmat->Info());
+  for (size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
+    auto row_idx = updaters_.back()->GetRowIndexCache(tree_idx);
+    // distinguish the difference between multi-class and multi-target.
+    auto target = targets > 1 ? gidx : 0;
+    obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), target, trees[tree_idx].get());
+  }
+}
+
 void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
                      PredictionCacheEntry* predt, ObjFunction const* obj) {
   std::vector<std::vector<std::unique_ptr<RegTree> > > new_trees;
@@ -235,11 +250,11 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
   if (ngroup == 1) {
     std::vector<std::unique_ptr<RegTree>> ret;
     BoostNewTrees(in_gpair, p_fmat, 0, &ret);
+    UpdateTreeLeaf(p_fmat, obj, 0, &ret);
     const size_t num_new_trees = ret.size();
     new_trees.push_back(std::move(ret));
     auto v_predt = out.Slice(linalg::All(), 0);
-    if (updaters_.size() > 0 && num_new_trees == 1 &&
-        predt->predictions.Size() > 0 &&
+    if (updaters_.size() > 0 && num_new_trees == 1 && predt->predictions.Size() > 0 &&
         updaters_.back()->UpdatePredictionCache(p_fmat, v_predt)) {
       predt->Update(1);
     }
@@ -254,6 +269,7 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
       CopyGradient(in_gpair, ctx_->Threads(), ngroup, gid, &tmp);
       std::vector<std::unique_ptr<RegTree> > ret;
       BoostNewTrees(&tmp, p_fmat, gid, &ret);
+      UpdateTreeLeaf(p_fmat, obj, gid, &ret);
       const size_t num_new_trees = ret.size();
       new_trees.push_back(std::move(ret));
       auto v_predt = out.Slice(linalg::All(), gid);
@@ -268,19 +284,6 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
     }
   }
 
-  if (obj) {
-    // Update tree leaf value at the end of boosting.
-    bst_group_t gidx{0};
-    auto targets = obj->Targets(p_fmat->Info());
-    for (auto& tree_group : new_trees) {
-      for (size_t t = 0; t < tree_group.size(); ++t) {
-        // within a forest
-        auto row_idx = updaters_.back()->GetRowIndexCache(t);
-        auto target = targets > 1 ? gidx : 0;
-        obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), gidx, tree_group[t].get());
-      }
-    }
-  }
   monitor_.Stop("BoostNewTrees");
   this->CommitModel(std::move(new_trees), p_fmat, predt);
 }
diff --git a/src/gbm/gbtree.h b/src/gbm/gbtree.h
index c973d0b02685..ea42c8f0457d 100644
--- a/src/gbm/gbtree.h
+++ b/src/gbm/gbtree.h
@@ -202,6 +202,9 @@ class GBTree : public GradientBooster {
   void ConfigureUpdaters();
   void ConfigureWithKnownData(Args const& cfg, DMatrix* fmat);
 
+  void UpdateTreeLeaf(DMatrix const* p_fmat, ObjFunction const* obj, size_t gidx,
+                      std::vector<std::unique_ptr<RegTree>>* p_trees);
+
   /*! \brief Carry out one iteration of boosting */
   void DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
                PredictionCacheEntry* predt, ObjFunction const* obj) override;

From 07523df5e72fe7c43385e13d7ce7b9ae21c323fd Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 17:27:51 +0800
Subject: [PATCH 010/124] Make sure it's used.

---
 include/xgboost/tree_model.h    |  3 +--
 src/common/row_set.h            |  8 ++++++++
 src/gbm/gbtree.cc               | 16 +++++++++-------
 src/objective/regression_obj.cc | 18 +++++++++++-------
 src/tree/hist/evaluate_splits.h |  6 ++++--
 src/tree/updater_approx.cc      | 12 ++++++++----
 6 files changed, 41 insertions(+), 22 deletions(-)

diff --git a/include/xgboost/tree_model.h b/include/xgboost/tree_model.h
index 8079b86d5c94..0f39c39c3d33 100644
--- a/include/xgboost/tree_model.h
+++ b/include/xgboost/tree_model.h
@@ -745,8 +745,7 @@ struct RowIndexCache {
   HostDeviceVector<size_t> row_index;
   std::vector<Segment> indptr;
 
-  RowIndexCache(Context const* ctx, size_t n_leaf, size_t n_samples) {
-    indptr.resize(n_leaf + 1);
+  RowIndexCache(Context const* ctx, size_t n_samples) {
     if (!ctx->IsCPU()) {
       row_index.SetDevice(ctx->gpu_id);
     }
diff --git a/src/common/row_set.h b/src/common/row_set.h
index 7261b02cc714..64f6089e8ec3 100644
--- a/src/common/row_set.h
+++ b/src/common/row_set.h
@@ -18,6 +18,12 @@ namespace common {
 /*! \brief collection of rowset */
 class RowSetCollection {
  public:
+  RowSetCollection() = default;
+  RowSetCollection(RowSetCollection const&) = delete;
+  RowSetCollection(RowSetCollection&&) = default;
+  RowSetCollection& operator=(RowSetCollection const&) = delete;
+  RowSetCollection& operator=(RowSetCollection&&) = default;
+
   /*! \brief data structure to store an instance set, a subset of
    *  rows (instances) associated with a particular node in a decision
    *  tree. */
@@ -82,6 +88,8 @@ class RowSetCollection {
   }
 
   std::vector<size_t>* Data() { return &row_indices_; }
+  std::vector<size_t> const* Data() const { return &row_indices_; }
+
   // split rowset into two
   inline void AddSplit(unsigned node_id, unsigned left_node_id, unsigned right_node_id,
                        size_t n_left, size_t n_right) {
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index f661e1bcbda4..a025820cf9ff 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -226,6 +226,11 @@ void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, ObjFunction const* obj, size_
   auto targets = obj->Targets(p_fmat->Info());
   for (size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
     auto row_idx = updaters_.back()->GetRowIndexCache(tree_idx);
+    for (auto const& part : row_idx) {
+      for (auto const& seg : part.indptr) {
+        CHECK((*trees[tree_idx])[seg.nidx].IsLeaf()) << "trees.size():" << trees.size();
+      }
+    }
     // distinguish the difference between multi-class and multi-target.
     auto target = targets > 1 ? gidx : 0;
     obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), target, trees[tree_idx].get());
@@ -259,22 +264,19 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
       predt->Update(1);
     }
   } else {
-    CHECK_EQ(in_gpair->Size() % ngroup, 0U)
-        << "must have exactly ngroup * nrow gpairs";
-    HostDeviceVector<GradientPair> tmp(in_gpair->Size() / ngroup,
-                                       GradientPair(),
+    CHECK_EQ(in_gpair->Size() % ngroup, 0U) << "must have exactly ngroup * nrow gpairs";
+    HostDeviceVector<GradientPair> tmp(in_gpair->Size() / ngroup, GradientPair(),
                                        in_gpair->DeviceIdx());
     bool update_predict = true;
     for (int gid = 0; gid < ngroup; ++gid) {
       CopyGradient(in_gpair, ctx_->Threads(), ngroup, gid, &tmp);
-      std::vector<std::unique_ptr<RegTree> > ret;
+      std::vector<std::unique_ptr<RegTree>> ret;
       BoostNewTrees(&tmp, p_fmat, gid, &ret);
       UpdateTreeLeaf(p_fmat, obj, gid, &ret);
       const size_t num_new_trees = ret.size();
       new_trees.push_back(std::move(ret));
       auto v_predt = out.Slice(linalg::All(), gid);
-      if (!(updaters_.size() > 0 && predt->predictions.Size() > 0 &&
-            num_new_trees == 1 &&
+      if (!(updaters_.size() > 0 && predt->predictions.Size() > 0 && num_new_trees == 1 &&
             updaters_.back()->UpdatePredictionCache(p_fmat, v_predt))) {
         update_predict = false;
       }
diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index 8f3ad23a27e0..8e1431fc3174 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -49,13 +49,14 @@ float WeightedQuantile(float quantile, common::Span<size_t const> row_set,
   }
 };
 
-void UpdateTreeLeafHost(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
-                        uint32_t target, float alpha, RegTree* p_tree) {
+void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
+                        MetaInfo const& info, uint32_t target, float alpha, RegTree* p_tree) {
   auto& tree = *p_tree;
   std::vector<float> quantiles;
   for (auto const& part : row_index) {
-    std::vector<float> results;
-    for (auto const& seg : part.indptr) {
+    std::vector<float> results(part.indptr.size());
+    common::ParallelFor(part.indptr.size(), ctx->Threads(), [&](size_t k) {
+      auto const& seg = part.indptr[k];
       auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
       float q{0};
       if (info.weights_.Empty()) {
@@ -65,8 +66,8 @@ void UpdateTreeLeafHost(common::Span<RowIndexCache const> row_index, MetaInfo co
         q = WeightedQuantile(alpha, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
                              linalg::MakeVec(&info.weights_));
       }
-      results.push_back(q);
-    }
+      results.at(k++) = q;
+    });
     // fixme: verify this is correct for external memory
     if (quantiles.empty()) {
       quantiles.resize(results.size(), 0);
@@ -85,7 +86,10 @@ void UpdateTreeLeafHost(common::Span<RowIndexCache const> row_index, MetaInfo co
   for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
     auto seg = row_index.front().indptr[i];
     auto q = quantiles[i];
+    auto l = tree[seg.nidx].LeafValue();
+    CHECK(tree[seg.nidx].IsLeaf());
     tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
+    l = tree[seg.nidx].LeafValue();
   }
 }
 
@@ -129,7 +133,7 @@ class MeanAbsoluteError : public ObjFunction {
 
   void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
                       uint32_t target, RegTree* p_tree) const override {
-    UpdateTreeLeafHost(row_index, info, target, 0.5, p_tree);
+    UpdateTreeLeafHost(ctx_, row_index, info, target, 0.5, p_tree);
   }
 
   const char* DefaultEvalMetric() const override { return "mae"; }
diff --git a/src/tree/hist/evaluate_splits.h b/src/tree/hist/evaluate_splits.h
index 053b485012bd..5646bfc85497 100644
--- a/src/tree/hist/evaluate_splits.h
+++ b/src/tree/hist/evaluate_splits.h
@@ -402,8 +402,10 @@ void UpdatePredictionCacheImpl(GenericParameter const *ctx, RegTree const *p_las
       if (!tree[nidx].IsDeleted() && tree[nidx].IsLeaf()) {
         auto const &rowset = part[nidx];
         auto const &stats = snode[nidx];
-        auto leaf_value =
-            evaluator.CalcWeight(nidx, param, GradStats{stats.stats}) * param.learning_rate;
+        auto leaf_value = tree[nidx].LeafValue();
+        // auto leaf_value =
+        //     evaluator.CalcWeight(nidx, param, GradStats{stats.stats}) * param.learning_rate;
+        // CHECK_EQ(leaf, leaf_value);
         for (const size_t *it = rowset.begin + r.begin(); it < rowset.begin + r.end(); ++it) {
           out_preds(*it) += leaf_value;
         }
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index d3f98baf0388..588c6c35b7a4 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -234,24 +234,28 @@ class GloablApproxBuilder {
 
     CHECK(p_out_row_indices->empty());
     for (auto const &part : partitioner_) {
-      auto row_set = part.Partitions();
-      auto n_leaf = row_set.Size();
-      p_out_row_indices->emplace_back(ctx_, n_leaf, row_set.Size());
+      auto const &row_set = part.Partitions();
+      p_out_row_indices->emplace_back(ctx_, p_fmat->Info().num_row_);
       // fixme: subsample
       auto &h_row_index = p_out_row_indices->back().row_index.HostVector();
 
       auto begin = row_set.Data()->data();
       for (auto node : row_set) {
-        CHECK(node.begin);
+        if (!node.begin) {
+          continue;
+        }
+        CHECK(node.begin)<< node.node_id;
         CHECK(tree[node.node_id].IsLeaf());
         size_t offset = node.begin - begin;
         auto size = node.Size();
+        CHECK_LT(offset, p_fmat->Info().num_row_) << node.node_id;
         auto seg = RowIndexCache::Segment{offset, size, node.node_id};
         size_t k = seg.begin;
         for (auto idx = node.begin; idx != node.end; ++idx) {
           h_row_index[k] = *idx;
         }
         p_out_row_indices->back().indptr.push_back(seg);
+        CHECK(tree[seg.nidx].IsLeaf());
       }
     }
   }

From fd1729cb89ed7b8cb58c4653464f1ba4cd3ea569 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 18:05:12 +0800
Subject: [PATCH 011/124] Use transform iter.

---
 include/xgboost/objective.h     |  3 +-
 src/common/stats.h              | 70 +++++++++++++++++++++++++++++----
 src/gbm/gbtree.cc               | 14 +++----
 src/gbm/gbtree.h                |  3 +-
 src/objective/regression_obj.cc | 18 ++++++---
 tests/cpp/common/test_stats.cc  | 12 +++---
 6 files changed, 91 insertions(+), 29 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 30ad4a2b0f71..a0dcd25c9795 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -101,7 +101,8 @@ class ObjFunction : public Configurable {
    * \param p_tree Tree that needs to be updated.
    */
   virtual void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
-                              uint32_t target, RegTree* p_tree) const {}
+                              HostDeviceVector<float> const& prediction, uint32_t target,
+                              RegTree* p_tree) const {}
 
   /*!
    * \brief Create an objective function according to name.
diff --git a/src/common/stats.h b/src/common/stats.h
index 884c51324f22..146d829931f2 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -3,6 +3,7 @@
  */
 #ifndef XGBOOST_COMMON_STATS_H_
 #define XGBOOST_COMMON_STATS_H_
+#include <iterator>
 #include <limits>
 #include <vector>
 
@@ -11,6 +12,59 @@
 
 namespace xgboost {
 namespace common {
+
+template <typename Fn>
+class IndexTransformIter {
+  size_t iter_{0};
+  Fn fn_;
+
+ public:
+  using iterator_category = std::random_access_iterator_tag;  // NOLINT
+  using value_type = std::result_of_t<Fn(size_t)>;            // NOLINT
+  using difference_type = detail::ptrdiff_t;                  // NOLINT
+  using reference = std::add_lvalue_reference_t<value_type>;  // NOLINT
+  using pointer = std::add_pointer_t<value_type>;             // NOLINT
+
+ public:
+  XGBOOST_DEVICE explicit IndexTransformIter(Fn&& fn) : fn_{fn} {}
+  XGBOOST_DEVICE IndexTransformIter(IndexTransformIter const&) = default;
+
+  XGBOOST_DEVICE value_type operator*() const { return fn_(iter_); }
+
+  XGBOOST_DEVICE auto operator-(IndexTransformIter const& that) const { return iter_ - that.iter_; }
+
+  XGBOOST_DEVICE IndexTransformIter& operator++() {
+    iter_++;
+    return *this;
+  }
+  XGBOOST_DEVICE IndexTransformIter operator++(int) {
+    auto ret = *this;
+    ++(*this);
+    return ret;
+  }
+  XGBOOST_DEVICE IndexTransformIter& operator+=(difference_type n) {
+    iter_ += n;
+    return *this;
+  }
+  XGBOOST_DEVICE IndexTransformIter& operator-=(difference_type n) {
+    (*this) += -n;
+    return *this;
+  }
+  XGBOOST_DEVICE IndexTransformIter operator+(difference_type n) const {
+    auto ret = *this;
+    return ret += n;
+  }
+  XGBOOST_DEVICE IndexTransformIter operator-(difference_type n) const {
+    auto ret = *this;
+    return ret -= n;
+  }
+};
+
+template <typename Fn>
+auto MakeIndexTransformIter(Fn&& fn) {
+  return IndexTransformIter<Fn>(std::forward<Fn>(fn));
+}
+
 /**
  * \brief Percentile with masked array using linear interpolation.
  *
@@ -22,19 +76,21 @@ namespace common {
  *
  * \return The result of interpolation.
  */
-inline float Percentile(double alpha, Span<size_t const> index,
-                        linalg::TensorView<float const, 1> arr) {
+template <typename Iter>
+float Percentile(double alpha, Iter begin, Iter end) {
   CHECK(alpha >= 0 && alpha <= 1);
-  if (index.size() == 0) {
+  auto n = static_cast<double>(std::distance(begin, end));
+  if (n == 0) {
     return std::numeric_limits<float>::quiet_NaN();
   }
-  auto n = static_cast<double>(index.size());
-  std::vector<size_t> sorted_idx(index.size());
+
+  std::vector<size_t> sorted_idx(n);
   std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
   std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
-                   [&](size_t l, size_t r) { return arr(index(l)) < arr(index(r)); });
+                   [&](size_t l, size_t r) { return *(begin + l) < *(begin + r); });
 
-  auto val = [&](size_t i) { return arr(index(sorted_idx[i])); };
+  auto val = [&](size_t i) { return *(begin + sorted_idx[i]); };
+  static_assert(std::is_same<decltype(val(0)), float>::value, "");
 
   if (alpha <= (1 / (n + 1))) {
     return val(0);
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index a025820cf9ff..d866705f9301 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -217,7 +217,8 @@ void CopyGradient(HostDeviceVector<GradientPair> const* in_gpair, int32_t n_thre
   }
 }
 
-void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, ObjFunction const* obj, size_t gidx,
+void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const& predictions,
+                            ObjFunction const* obj, size_t gidx,
                             std::vector<std::unique_ptr<RegTree>>* p_trees) {
   if (!obj) {
     return;
@@ -226,14 +227,9 @@ void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, ObjFunction const* obj, size_
   auto targets = obj->Targets(p_fmat->Info());
   for (size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
     auto row_idx = updaters_.back()->GetRowIndexCache(tree_idx);
-    for (auto const& part : row_idx) {
-      for (auto const& seg : part.indptr) {
-        CHECK((*trees[tree_idx])[seg.nidx].IsLeaf()) << "trees.size():" << trees.size();
-      }
-    }
     // distinguish the difference between multi-class and multi-target.
     auto target = targets > 1 ? gidx : 0;
-    obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), target, trees[tree_idx].get());
+    obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), predictions, target, trees[tree_idx].get());
   }
 }
 
@@ -255,7 +251,7 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
   if (ngroup == 1) {
     std::vector<std::unique_ptr<RegTree>> ret;
     BoostNewTrees(in_gpair, p_fmat, 0, &ret);
-    UpdateTreeLeaf(p_fmat, obj, 0, &ret);
+    UpdateTreeLeaf(p_fmat, predt->predictions, obj, 0, &ret);
     const size_t num_new_trees = ret.size();
     new_trees.push_back(std::move(ret));
     auto v_predt = out.Slice(linalg::All(), 0);
@@ -272,7 +268,7 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
       CopyGradient(in_gpair, ctx_->Threads(), ngroup, gid, &tmp);
       std::vector<std::unique_ptr<RegTree>> ret;
       BoostNewTrees(&tmp, p_fmat, gid, &ret);
-      UpdateTreeLeaf(p_fmat, obj, gid, &ret);
+      UpdateTreeLeaf(p_fmat, predt->predictions, obj, gid, &ret);
       const size_t num_new_trees = ret.size();
       new_trees.push_back(std::move(ret));
       auto v_predt = out.Slice(linalg::All(), gid);
diff --git a/src/gbm/gbtree.h b/src/gbm/gbtree.h
index ea42c8f0457d..0e5343cec5c7 100644
--- a/src/gbm/gbtree.h
+++ b/src/gbm/gbtree.h
@@ -202,7 +202,8 @@ class GBTree : public GradientBooster {
   void ConfigureUpdaters();
   void ConfigureWithKnownData(Args const& cfg, DMatrix* fmat);
 
-  void UpdateTreeLeaf(DMatrix const* p_fmat, ObjFunction const* obj, size_t gidx,
+  void UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const& predictions,
+                      ObjFunction const* obj, size_t gidx,
                       std::vector<std::unique_ptr<RegTree>>* p_trees);
 
   /*! \brief Carry out one iteration of boosting */
diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index 8e1431fc3174..e9bcfdb7e8eb 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -50,7 +50,8 @@ float WeightedQuantile(float quantile, common::Span<size_t const> row_set,
 };
 
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
-                        MetaInfo const& info, uint32_t target, float alpha, RegTree* p_tree) {
+                        MetaInfo const& info, HostDeviceVector<float> const& prediction,
+                        uint32_t target, float alpha, RegTree* p_tree) {
   auto& tree = *p_tree;
   std::vector<float> quantiles;
   for (auto const& part : row_index) {
@@ -59,9 +60,15 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
       auto const& seg = part.indptr[k];
       auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
       float q{0};
+      auto h_labels = info.labels.HostView().Slice(linalg::All(), target);
+      auto const& h_prediction = prediction.ConstHostVector();
+      auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
+        auto row_idx = h_row_set[i];
+        return h_labels(h_row_set[i]) - h_prediction[row_idx];
+      });
+
       if (info.weights_.Empty()) {
-        q = common::Percentile(alpha, h_row_set,
-                               info.labels.HostView().Slice(linalg::All(), target));
+        q = common::Percentile(alpha, iter, iter + h_row_set.size());
       } else {
         q = WeightedQuantile(alpha, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
                              linalg::MakeVec(&info.weights_));
@@ -132,8 +139,9 @@ class MeanAbsoluteError : public ObjFunction {
   }
 
   void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
-                      uint32_t target, RegTree* p_tree) const override {
-    UpdateTreeLeafHost(ctx_, row_index, info, target, 0.5, p_tree);
+                      HostDeviceVector<float> const& prediction, uint32_t target,
+                      RegTree* p_tree) const override {
+    UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
   }
 
   const char* DefaultEvalMetric() const override { return "mae"; }
diff --git a/tests/cpp/common/test_stats.cc b/tests/cpp/common/test_stats.cc
index 49ab0fb98911..327c709bde67 100644
--- a/tests/cpp/common/test_stats.cc
+++ b/tests/cpp/common/test_stats.cc
@@ -8,14 +8,14 @@ namespace common {
 TEST(Stats, Percentil) {
   linalg::Tensor<float, 1> arr({20.f, 0.f, 15.f, 50.f, 40.f, 0.f, 35.f}, {7}, Context::kCpuId);
   std::vector<size_t> index{0, 2, 3, 4, 6};
-  auto percentile = Percentile(0.40f, Span<size_t const>{index}, arr.HostView());
-  ASSERT_EQ(percentile, 26.0);
+  // auto percentile = Percentile(0.40f, Span<size_t const>{index}, arr.HostView());
+  // ASSERT_EQ(percentile, 26.0);
 
-  percentile = Percentile(0.20f, Span<size_t const>{index}, arr.HostView());
-  ASSERT_EQ(percentile, 16.0);
+  // percentile = Percentile(0.20f, Span<size_t const>{index}, arr.HostView());
+  // ASSERT_EQ(percentile, 16.0);
 
-  percentile = Percentile(0.10f, Span<size_t const>{index}, arr.HostView());
-  ASSERT_EQ(percentile, 15.0);
+  // percentile = Percentile(0.10f, Span<size_t const>{index}, arr.HostView());
+  // ASSERT_EQ(percentile, 15.0);
 }
 }  // namespace common
 }  // namespace xgboost

From 033ac6f89e6c58a19c304e80bd675297b00eee0b Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 18:18:18 +0800
Subject: [PATCH 012/124] fixes.

---
 src/common/stats.h              |  8 ++++----
 src/objective/regression_obj.cc | 17 ++++++++++-------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/common/stats.h b/src/common/stats.h
index 146d829931f2..d8b976ff0e1f 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -70,14 +70,14 @@ auto MakeIndexTransformIter(Fn&& fn) {
  *
  *   https://www.itl.nist.gov/div898/handbook/prc/section2/prc262.htm
  *
- * \param alpha percentile, must be in range [0, 1].
- * \param index The index of valid elements in arr.
- * \param arr   Input values.
+ * \param alpha Percentile, must be in range [0, 1].
+ * \param begin Iterator begin for input array.
+ * \param end   Iterator end for input array.
  *
  * \return The result of interpolation.
  */
 template <typename Iter>
-float Percentile(double alpha, Iter begin, Iter end) {
+float Percentile(double alpha, Iter const& begin, Iter const& end) {
   CHECK(alpha >= 0 && alpha <= 1);
   auto n = static_cast<double>(std::distance(begin, end));
   if (n == 0) {
diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index e9bcfdb7e8eb..ba14a68a93f1 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -58,13 +58,14 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     std::vector<float> results(part.indptr.size());
     common::ParallelFor(part.indptr.size(), ctx->Threads(), [&](size_t k) {
       auto const& seg = part.indptr[k];
+      CHECK(tree[seg.nidx].IsLeaf());
       auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
       float q{0};
       auto h_labels = info.labels.HostView().Slice(linalg::All(), target);
       auto const& h_prediction = prediction.ConstHostVector();
       auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
         auto row_idx = h_row_set[i];
-        return h_labels(h_row_set[i]) - h_prediction[row_idx];
+        return h_labels(row_idx) - h_prediction[row_idx];
       });
 
       if (info.weights_.Empty()) {
@@ -73,8 +74,9 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
         q = WeightedQuantile(alpha, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
                              linalg::MakeVec(&info.weights_));
       }
-      results.at(k++) = q;
+      results.at(k) = q;
     });
+
     // fixme: verify this is correct for external memory
     if (quantiles.empty()) {
       quantiles.resize(results.size(), 0);
@@ -82,13 +84,14 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     for (size_t i = 0; i < results.size(); ++i) {
       quantiles[i] += results[i];
     }
-    // use the mean value
-    rabit::Allreduce<rabit::op::Sum>(results.data(), results.size());
-    auto world = rabit::GetWorldSize();
-    std::transform(results.begin(), results.end(), results.begin(),
-                   [&](float q) { return q / world; });
   }
 
+  // use the mean value
+  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
+  auto world = rabit::GetWorldSize();
+  std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
+                 [&](float q) { return q / world; });
+
   // fixme: verify this is correct for external memory
   for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
     auto seg = row_index.front().indptr[i];

From 6b53665cabd14c605e6e8b9815e0609759bad4c2 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 18:45:11 +0800
Subject: [PATCH 013/124] Fix.

---
 src/tree/updater_approx.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index 588c6c35b7a4..213a1b068880 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -252,7 +252,7 @@ class GloablApproxBuilder {
         auto seg = RowIndexCache::Segment{offset, size, node.node_id};
         size_t k = seg.begin;
         for (auto idx = node.begin; idx != node.end; ++idx) {
-          h_row_index[k] = *idx;
+          h_row_index[k++] = *idx;
         }
         p_out_row_indices->back().indptr.push_back(seg);
         CHECK(tree[seg.nidx].IsLeaf());

From dc1015c81187b348a3fd7f109d1f08b862d6486e Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 18:48:37 +0800
Subject: [PATCH 014/124] cleanup.

---
 src/tree/updater_approx.cc | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index 213a1b068880..4070bfe98546 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -244,8 +244,7 @@ class GloablApproxBuilder {
         if (!node.begin) {
           continue;
         }
-        CHECK(node.begin)<< node.node_id;
-        CHECK(tree[node.node_id].IsLeaf());
+        CHECK(node.begin && tree[node.node_id].IsLeaf()) << " Offending node idx:" << node.node_id;
         size_t offset = node.begin - begin;
         auto size = node.Size();
         CHECK_LT(offset, p_fmat->Info().num_row_) << node.node_id;
@@ -255,7 +254,6 @@ class GloablApproxBuilder {
           h_row_index[k++] = *idx;
         }
         p_out_row_indices->back().indptr.push_back(seg);
-        CHECK(tree[seg.nidx].IsLeaf());
       }
     }
   }

From 076f810e15580c35d8229a380433a95c13d335f0 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 18:54:58 +0800
Subject: [PATCH 015/124] Subsample.

---
 src/tree/updater_approx.cc | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index 4070bfe98546..d51df8760fb8 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -236,7 +236,6 @@ class GloablApproxBuilder {
     for (auto const &part : partitioner_) {
       auto const &row_set = part.Partitions();
       p_out_row_indices->emplace_back(ctx_, p_fmat->Info().num_row_);
-      // fixme: subsample
       auto &h_row_index = p_out_row_indices->back().row_index.HostVector();
 
       auto begin = row_set.Data()->data();
@@ -246,13 +245,14 @@ class GloablApproxBuilder {
         }
         CHECK(node.begin && tree[node.node_id].IsLeaf()) << " Offending node idx:" << node.node_id;
         size_t offset = node.begin - begin;
-        auto size = node.Size();
         CHECK_LT(offset, p_fmat->Info().num_row_) << node.node_id;
-        auto seg = RowIndexCache::Segment{offset, size, node.node_id};
-        size_t k = seg.begin;
+        size_t k = offset;
         for (auto idx = node.begin; idx != node.end; ++idx) {
-          h_row_index[k++] = *idx;
+          if (hess[*idx] != 0.f) {
+            h_row_index[k++] = *idx;
+          }
         }
+        auto seg = RowIndexCache::Segment{offset, k - offset, node.node_id};
         p_out_row_indices->back().indptr.push_back(seg);
       }
     }

From 609ceb347827dfa6420e3afff6c311b6a4c27976 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 19:22:35 +0800
Subject: [PATCH 016/124] Start working on GPU.

---
 src/objective/regression_obj.cc |  2 --
 src/objective/regression_obj.cu | 15 +++++++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index ba14a68a93f1..0f6611f3106e 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -96,10 +96,8 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
   for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
     auto seg = row_index.front().indptr[i];
     auto q = quantiles[i];
-    auto l = tree[seg.nidx].LeafValue();
     CHECK(tree[seg.nidx].IsLeaf());
     tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
-    l = tree[seg.nidx].LeafValue();
   }
 }
 
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index fa294a5a5773..725658fa2ffb 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -19,6 +19,7 @@
 #include "../common/threading_utils.h"
 #include "../common/transform.h"
 #include "./regression_loss.h"
+#include "xgboost/generic_parameters.h"
 #include "xgboost/host_device_vector.h"
 #include "xgboost/json.h"
 #include "xgboost/parameter.h"
@@ -675,5 +676,19 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .describe("Tweedie regression for insurance data.")
 .set_body([]() { return new TweedieRegression(); });
 
+void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const> row_index,
+                          MetaInfo const& info, HostDeviceVector<float> const& prediction,
+                          uint32_t target, float alpha, RegTree* p_tree) {
+  dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
+  CHECK_EQ(row_index.size(), 1);
+  auto part = row_index.front();
+
+  dh::caching_device_vector<float> quantiles;
+  auto d_quantiles = dh::ToSpan(quantiles);
+  dh::LaunchN(part.row_index.Size(), [=]XGBOOST_DEVICE(size_t i) {
+
+  });
+}
+
 }  // namespace obj
 }  // namespace xgboost

From f5732dcbe6effe988c1a3eaccb8f22610fc80d30 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 8 Apr 2022 22:53:01 +0800
Subject: [PATCH 017/124] Work on GPU.

---
 include/xgboost/task.h          |   6 +-
 plugin/example/custom_obj.cc    |   8 +-
 src/common/stats.h              |  35 +++++++-
 src/objective/aft_obj.cu        |   4 +-
 src/objective/hinge.cu          |   6 +-
 src/objective/multiclass_obj.cu |   2 +-
 src/objective/rank_obj.cu       |   7 +-
 src/objective/regression_loss.h |  10 +--
 src/objective/regression_obj.cc | 142 +------------------------------
 src/objective/regression_obj.cu | 145 +++++++++++++++++++++++++++-----
 src/tree/updater_gpu_hist.cu    |  96 ++++++++++++++-------
 11 files changed, 240 insertions(+), 221 deletions(-)

diff --git a/include/xgboost/task.h b/include/xgboost/task.h
index 537320657544..91b02a52e438 100644
--- a/include/xgboost/task.h
+++ b/include/xgboost/task.h
@@ -7,6 +7,7 @@
 #include <xgboost/base.h>
 
 #include <cinttypes>
+#include <cstdint>
 
 namespace xgboost {
 /*!
@@ -33,9 +34,10 @@ struct ObjInfo {
   } task;
   // Does the objective have constant hessian value?
   bool const_hess{false};
+  bool zero_hess{false};
 
-  explicit ObjInfo(Task t) : task{t} {}
-  ObjInfo(Task t, bool khess) : task{t}, const_hess{khess} {}
+  ObjInfo(Task t) : task{t} {}  // NOLINT
+  ObjInfo(Task t, bool khess, bool zhess) : task{t}, const_hess{khess}, zero_hess(zhess) {}
 
   XGBOOST_DEVICE bool UseOneHot() const {
     return (task != ObjInfo::kRegression && task != ObjInfo::kBinary);
diff --git a/plugin/example/custom_obj.cc b/plugin/example/custom_obj.cc
index b61073360e00..1b26fea410a6 100644
--- a/plugin/example/custom_obj.cc
+++ b/plugin/example/custom_obj.cc
@@ -31,13 +31,9 @@ DMLC_REGISTER_PARAMETER(MyLogisticParam);
 // Implement the interface.
 class MyLogistic : public ObjFunction {
  public:
-  void Configure(const std::vector<std::pair<std::string, std::string> >& args) override {
-    param_.UpdateAllowUnknown(args);
-  }
+  void Configure(const Args& args) override { param_.UpdateAllowUnknown(args); }
 
-  struct ObjInfo Task() const override {
-    return {ObjInfo::kRegression, false};
-  }
+  ObjInfo Task() const override { return ObjInfo::kRegression; }
 
   void GetGradient(const HostDeviceVector<bst_float> &preds,
                    const MetaInfo &info,
diff --git a/src/common/stats.h b/src/common/stats.h
index d8b976ff0e1f..dfb6feff2178 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -27,9 +27,9 @@ class IndexTransformIter {
 
  public:
   XGBOOST_DEVICE explicit IndexTransformIter(Fn&& fn) : fn_{fn} {}
-  XGBOOST_DEVICE IndexTransformIter(IndexTransformIter const&) = default;
+  IndexTransformIter(IndexTransformIter const&) = default;
 
-  XGBOOST_DEVICE value_type operator*() const { return fn_(iter_); }
+  value_type operator*() const { return fn_(iter_); }
 
   XGBOOST_DEVICE auto operator-(IndexTransformIter const& that) const { return iter_ - that.iter_; }
 
@@ -108,6 +108,37 @@ float Percentile(double alpha, Iter const& begin, Iter const& end) {
   auto v1 = val(static_cast<size_t>(k) + 1);
   return v0 + d * (v1 - v0);
 }
+
+inline float WeightedPercentile(float quantile, common::Span<size_t const> row_set,
+                                linalg::VectorView<float const> labels,
+                                linalg::VectorView<float const> weights) {
+  std::vector<size_t> sorted_idx(row_set.size());
+  std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
+  std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
+                   [&](size_t i, size_t j) { return labels(row_set[i]) < labels(row_set[j]); });
+  std::vector<float> weighted_cdf(row_set.size());
+  weighted_cdf[0] = weights(row_set[sorted_idx[0]]);
+  for (size_t i = 1; i < row_set.size(); ++i) {
+    weighted_cdf[i] = weighted_cdf[i - 1] + weights(row_set[sorted_idx[i]]);
+  }
+  float thresh = weighted_cdf.back() * quantile;
+  size_t pos =
+      std::upper_bound(weighted_cdf.cbegin(), weighted_cdf.cend(), thresh) - weighted_cdf.cbegin();
+  pos = std::min(pos, static_cast<size_t>(row_set.size() - 1));
+  if (pos == 0 || pos == static_cast<size_t>(row_set.size() - 1)) {
+    return labels(row_set[sorted_idx[pos]]);
+  }
+  CHECK_GE(thresh, weighted_cdf[pos - 1]);
+  CHECK_LT(thresh, weighted_cdf[pos]);
+  float v1 = labels(row_set[sorted_idx[pos - 1]]);
+  float v2 = labels(row_set[sorted_idx[pos]]);
+  if (weighted_cdf[pos + 1] - weighted_cdf[pos] >= 1.0f) {
+    return (thresh - weighted_cdf[pos]) / (weighted_cdf[pos + 1] - weighted_cdf[pos]) * (v2 - v2) +
+           v1;
+  } else {
+    return v2;
+  }
+}
 }  // namespace common
 }  // namespace xgboost
 #endif  // XGBOOST_COMMON_STATS_H_
diff --git a/src/objective/aft_obj.cu b/src/objective/aft_obj.cu
index 0e2d9290f95c..5f2306dee082 100644
--- a/src/objective/aft_obj.cu
+++ b/src/objective/aft_obj.cu
@@ -34,11 +34,11 @@ DMLC_REGISTRY_FILE_TAG(aft_obj_gpu);
 
 class AFTObj : public ObjFunction {
  public:
-  void Configure(const std::vector<std::pair<std::string, std::string> >& args) override {
+  void Configure(Args const& args) override {
     param_.UpdateAllowUnknown(args);
   }
 
-  ObjInfo Task() const override { return {ObjInfo::kSurvival, false}; }
+  ObjInfo Task() const override { return ObjInfo::kSurvival; }
 
   template <typename Distribution>
   void GetGradientImpl(const HostDeviceVector<bst_float> &preds,
diff --git a/src/objective/hinge.cu b/src/objective/hinge.cu
index e1f0df74d4e1..e062b2b48e3c 100644
--- a/src/objective/hinge.cu
+++ b/src/objective/hinge.cu
@@ -24,10 +24,8 @@ class HingeObj : public ObjFunction {
  public:
   HingeObj() = default;
 
-  void Configure(
-      const std::vector<std::pair<std::string, std::string> > &args) override {}
-
-  ObjInfo Task() const override { return {ObjInfo::kRegression, false}; }
+  void Configure(Args const&) override {}
+  ObjInfo Task() const override { return ObjInfo::kRegression; }
 
   void GetGradient(const HostDeviceVector<bst_float> &preds,
                    const MetaInfo &info,
diff --git a/src/objective/multiclass_obj.cu b/src/objective/multiclass_obj.cu
index 4b912a81710d..312992ec59f2 100644
--- a/src/objective/multiclass_obj.cu
+++ b/src/objective/multiclass_obj.cu
@@ -46,7 +46,7 @@ class SoftmaxMultiClassObj : public ObjFunction {
     param_.UpdateAllowUnknown(args);
   }
 
-  ObjInfo Task() const override { return {ObjInfo::kClassification, false}; }
+  ObjInfo Task() const override { return ObjInfo::kClassification; }
 
   void GetGradient(const HostDeviceVector<bst_float>& preds,
                    const MetaInfo& info,
diff --git a/src/objective/rank_obj.cu b/src/objective/rank_obj.cu
index 0bbf6f6df26b..899529824320 100644
--- a/src/objective/rank_obj.cu
+++ b/src/objective/rank_obj.cu
@@ -750,11 +750,8 @@ class SortedLabelList : dh::SegmentSorter<float> {
 template <typename LambdaWeightComputerT>
 class LambdaRankObj : public ObjFunction {
  public:
-  void Configure(const std::vector<std::pair<std::string, std::string> >& args) override {
-    param_.UpdateAllowUnknown(args);
-  }
-
-  ObjInfo Task() const override { return {ObjInfo::kRanking, false}; }
+  void Configure(Args const &args) override { param_.UpdateAllowUnknown(args); }
+  ObjInfo Task() const override { return ObjInfo::kRanking; }
 
   void GetGradient(const HostDeviceVector<bst_float>& preds,
                    const MetaInfo& info,
diff --git a/src/objective/regression_loss.h b/src/objective/regression_loss.h
index f92dfe2d47d7..8d9d661f23ef 100644
--- a/src/objective/regression_loss.h
+++ b/src/objective/regression_loss.h
@@ -38,7 +38,7 @@ struct LinearSquareLoss {
   static const char* DefaultEvalMetric() { return "rmse"; }
 
   static const char* Name() { return "reg:squarederror"; }
-  static ObjInfo Info() { return {ObjInfo::kRegression, true}; }
+  static ObjInfo Info() { return ObjInfo::kRegression; }
 };
 
 struct SquaredLogError {
@@ -65,7 +65,7 @@ struct SquaredLogError {
 
   static const char* Name() { return "reg:squaredlogerror"; }
 
-  static ObjInfo Info() { return {ObjInfo::kRegression, false}; }
+  static ObjInfo Info() { return ObjInfo::kRegression; }
 };
 
 // logistic loss for probability regression task
@@ -102,14 +102,14 @@ struct LogisticRegression {
 
   static const char* Name() { return "reg:logistic"; }
 
-  static ObjInfo Info() { return {ObjInfo::kRegression, false}; }
+  static ObjInfo Info() { return ObjInfo::kRegression; }
 };
 
 // logistic loss for binary classification task
 struct LogisticClassification : public LogisticRegression {
   static const char* DefaultEvalMetric() { return "logloss"; }
   static const char* Name() { return "binary:logistic"; }
-  static ObjInfo Info() { return {ObjInfo::kBinary, false}; }
+  static ObjInfo Info() { return ObjInfo::kBinary; }
 };
 
 // logistic loss, but predict un-transformed margin
@@ -146,7 +146,7 @@ struct LogisticRaw : public LogisticRegression {
 
   static const char* Name() { return "binary:logitraw"; }
 
-  static ObjInfo Info() { return {ObjInfo::kRegression, false}; }
+  static ObjInfo Info() { return ObjInfo::kRegression; }
 };
 
 }  // namespace obj
diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index 0f6611f3106e..0a80064850fd 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -7,7 +7,7 @@
 #include <dmlc/registry.h>
 
 #include "../common/linalg_op.h"
-#include "../common/stats.h"
+
 #include "rabit/rabit.h"
 #include "xgboost/data.h"
 #include "xgboost/objective.h"
@@ -18,146 +18,6 @@ namespace obj {
 
 DMLC_REGISTRY_FILE_TAG(regression_obj);
 
-float WeightedQuantile(float quantile, common::Span<size_t const> row_set,
-                       linalg::VectorView<float const> labels,
-                       linalg::VectorView<float const> weights) {
-  std::vector<size_t> sorted_idx(row_set.size());
-  std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
-  std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
-                   [&](size_t i, size_t j) { return labels(row_set[i]) < labels(row_set[j]); });
-  std::vector<float> weighted_cdf(row_set.size());
-  weighted_cdf[0] = weights(row_set[sorted_idx[0]]);
-  for (size_t i = 1; i < row_set.size(); ++i) {
-    weighted_cdf[i] = weighted_cdf[i - 1] + weights(row_set[sorted_idx[i]]);
-  }
-  float thresh = weighted_cdf.back() * quantile;
-  size_t pos =
-      std::upper_bound(weighted_cdf.cbegin(), weighted_cdf.cend(), thresh) - weighted_cdf.cbegin();
-  pos = std::min(pos, static_cast<size_t>(row_set.size() - 1));
-  if (pos == 0 || pos == static_cast<size_t>(row_set.size() - 1)) {
-    return labels(row_set[sorted_idx[pos]]);
-  }
-  CHECK_GE(thresh, weighted_cdf[pos - 1]);
-  CHECK_LT(thresh, weighted_cdf[pos]);
-  float v1 = labels(row_set[sorted_idx[pos - 1]]);
-  float v2 = labels(row_set[sorted_idx[pos]]);
-  if (weighted_cdf[pos + 1] - weighted_cdf[pos] >= 1.0f) {
-    return (thresh - weighted_cdf[pos]) / (weighted_cdf[pos + 1] - weighted_cdf[pos]) * (v2 - v2) +
-           v1;
-  } else {
-    return v2;
-  }
-};
-
-void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
-                        MetaInfo const& info, HostDeviceVector<float> const& prediction,
-                        uint32_t target, float alpha, RegTree* p_tree) {
-  auto& tree = *p_tree;
-  std::vector<float> quantiles;
-  for (auto const& part : row_index) {
-    std::vector<float> results(part.indptr.size());
-    common::ParallelFor(part.indptr.size(), ctx->Threads(), [&](size_t k) {
-      auto const& seg = part.indptr[k];
-      CHECK(tree[seg.nidx].IsLeaf());
-      auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
-      float q{0};
-      auto h_labels = info.labels.HostView().Slice(linalg::All(), target);
-      auto const& h_prediction = prediction.ConstHostVector();
-      auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
-        auto row_idx = h_row_set[i];
-        return h_labels(row_idx) - h_prediction[row_idx];
-      });
-
-      if (info.weights_.Empty()) {
-        q = common::Percentile(alpha, iter, iter + h_row_set.size());
-      } else {
-        q = WeightedQuantile(alpha, h_row_set, info.labels.HostView().Slice(linalg::All(), target),
-                             linalg::MakeVec(&info.weights_));
-      }
-      results.at(k) = q;
-    });
-
-    // fixme: verify this is correct for external memory
-    if (quantiles.empty()) {
-      quantiles.resize(results.size(), 0);
-    }
-    for (size_t i = 0; i < results.size(); ++i) {
-      quantiles[i] += results[i];
-    }
-  }
-
-  // use the mean value
-  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
-  auto world = rabit::GetWorldSize();
-  std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
-                 [&](float q) { return q / world; });
-
-  // fixme: verify this is correct for external memory
-  for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
-    auto seg = row_index.front().indptr[i];
-    auto q = quantiles[i];
-    CHECK(tree[seg.nidx].IsLeaf());
-    tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
-  }
-}
-
-class MeanAbsoluteError : public ObjFunction {
- public:
-  void Configure(Args const&) override {}
-
-  uint32_t Targets(MetaInfo const& info) const override {
-    return std::max(static_cast<size_t>(1), info.labels.Shape(1));
-  }
-
-  struct ObjInfo Task() const override {
-    return {ObjInfo::kRegression, true};
-  }
-
-  void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
-                   HostDeviceVector<GradientPair>* out_gpair) override {
-    auto labels = info.labels.View(ctx_->gpu_id);
-
-    out_gpair->SetDevice(ctx_->gpu_id);
-    out_gpair->Resize(info.labels.Size());
-    auto gpair = linalg::MakeVec(out_gpair);
-
-    preds.SetDevice(ctx_->gpu_id);
-    auto predt = linalg::MakeVec(&preds);
-    auto sign = [](auto x) {
-      return (x > static_cast<decltype(x)>(0)) - (x < static_cast<decltype(x)>(0));
-    };
-
-    info.weights_.SetDevice(ctx_->gpu_id);
-    common::OptionalWeights weight{ctx_->IsCPU() ? info.weights_.ConstHostSpan()
-                                                 : info.weights_.ConstDeviceSpan()};
-
-    linalg::ElementWiseKernel(ctx_, labels, [=] XGBOOST_DEVICE(size_t i, float const y) mutable {
-      auto sample_id = std::get<0>(linalg::UnravelIndex(i, labels.Shape()));
-      auto grad = sign(predt(i) - y) * weight[i];
-      auto hess = weight[sample_id];
-      gpair(i) = GradientPair{grad, hess};
-    });
-  }
-
-  void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
-                      HostDeviceVector<float> const& prediction, uint32_t target,
-                      RegTree* p_tree) const override {
-    UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
-  }
-
-  const char* DefaultEvalMetric() const override { return "mae"; }
-
-  void SaveConfig(Json* p_out) const override {
-    auto& out = *p_out;
-    out["name"] = String("reg:absoluteerror");
-  }
-
-  void LoadConfig(Json const& in) override {}
-};
-
-XGBOOST_REGISTER_OBJECTIVE(MeanAbsoluteError, "reg:absoluteerror")
-    .describe("Mean absoluate error.")
-    .set_body([]() { return new MeanAbsoluteError(); });
 }  // namespace obj
 }  // namespace xgboost
 
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 725658fa2ffb..324a8d166b74 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -8,6 +8,7 @@
 #include <dmlc/omp.h>
 #include <xgboost/logging.h>
 #include <xgboost/objective.h>
+#include <xgboost/tree_model.h>
 
 #include <cmath>
 #include <memory>
@@ -16,6 +17,7 @@
 #include "../common/common.h"
 #include "../common/linalg_op.h"
 #include "../common/pseudo_huber.h"
+#include "../common/stats.h"
 #include "../common/threading_utils.h"
 #include "../common/transform.h"
 #include "./regression_loss.h"
@@ -68,9 +70,7 @@ class RegLossObj : public ObjFunction {
     param_.UpdateAllowUnknown(args);
   }
 
-  struct ObjInfo Task() const override {
-    return Loss::Info();
-  }
+  ObjInfo Task() const override { return Loss::Info(); }
 
   uint32_t Targets(MetaInfo const& info) const override {
     // Multi-target regression.
@@ -210,7 +210,7 @@ class PseudoHuberRegression : public ObjFunction {
 
  public:
   void Configure(Args const& args) override { param_.UpdateAllowUnknown(args); }
-  struct ObjInfo Task() const override { return {ObjInfo::kRegression, false}; }
+  ObjInfo Task() const override { return ObjInfo::kRegression; }
   uint32_t Targets(MetaInfo const& info) const override {
     return std::max(static_cast<size_t>(1), info.labels.Shape(1));
   }
@@ -287,9 +287,7 @@ class PoissonRegression : public ObjFunction {
     param_.UpdateAllowUnknown(args);
   }
 
-  struct ObjInfo Task() const override {
-    return {ObjInfo::kRegression, false};
-  }
+  ObjInfo Task() const override { return ObjInfo::kRegression; }
 
   void GetGradient(const HostDeviceVector<bst_float>& preds,
                    const MetaInfo &info, int,
@@ -379,12 +377,8 @@ XGBOOST_REGISTER_OBJECTIVE(PoissonRegression, "count:poisson")
 // cox regression for survival data (negative values mean they are censored)
 class CoxRegression : public ObjFunction {
  public:
-  void Configure(
-      const std::vector<std::pair<std::string, std::string> >&) override {}
-
-  struct ObjInfo Task() const override {
-    return {ObjInfo::kRegression, false};
-  }
+  void Configure(Args const&) override {}
+  ObjInfo Task() const override { return ObjInfo::kRegression; }
 
   void GetGradient(const HostDeviceVector<bst_float>& preds,
                    const MetaInfo &info, int,
@@ -480,12 +474,8 @@ XGBOOST_REGISTER_OBJECTIVE(CoxRegression, "survival:cox")
 // gamma regression
 class GammaRegression : public ObjFunction {
  public:
-  void Configure(
-      const std::vector<std::pair<std::string, std::string> >&) override {}
-
-  struct ObjInfo Task() const override {
-    return {ObjInfo::kRegression, false};
-  }
+  void Configure(Args const&) override {}
+  ObjInfo Task() const override { return ObjInfo::kRegression; }
 
   void GetGradient(const HostDeviceVector<bst_float> &preds,
                    const MetaInfo &info, int,
@@ -583,9 +573,7 @@ class TweedieRegression : public ObjFunction {
     metric_ = os.str();
   }
 
-  struct ObjInfo Task() const override {
-    return {ObjInfo::kRegression, false};
-  }
+  ObjInfo Task() const override { return ObjInfo::kRegression; }
 
   void GetGradient(const HostDeviceVector<bst_float>& preds,
                    const MetaInfo &info, int,
@@ -681,14 +669,125 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
                           uint32_t target, float alpha, RegTree* p_tree) {
   dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
   CHECK_EQ(row_index.size(), 1);
-  auto part = row_index.front();
+  auto const& part = row_index.front();
 
   dh::caching_device_vector<float> quantiles;
   auto d_quantiles = dh::ToSpan(quantiles);
+
+  auto rows = part.row_index.ConstDeviceSpan();
   dh::LaunchN(part.row_index.Size(), [=]XGBOOST_DEVICE(size_t i) {
 
   });
 }
 
+void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
+                        MetaInfo const& info, HostDeviceVector<float> const& prediction,
+                        uint32_t target, float alpha, RegTree* p_tree) {
+  auto& tree = *p_tree;
+  std::vector<float> quantiles;
+  for (auto const& part : row_index) {
+    std::vector<float> results(part.indptr.size());
+    common::ParallelFor(part.indptr.size(), ctx->Threads(), [&](size_t k) {
+      auto const& seg = part.indptr[k];
+      CHECK(tree[seg.nidx].IsLeaf());
+      auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
+      float q{0};
+      auto h_labels = info.labels.HostView().Slice(linalg::All(), target);
+      auto const& h_prediction = prediction.ConstHostVector();
+      auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
+        auto row_idx = h_row_set[i];
+        return h_labels(row_idx) - h_prediction[row_idx];
+      });
+
+      if (info.weights_.Empty()) {
+        q = common::Percentile(alpha, iter, iter + h_row_set.size());
+      } else {
+        q = common::WeightedPercentile(alpha, h_row_set,
+                                       info.labels.HostView().Slice(linalg::All(), target),
+                                       linalg::MakeVec(&info.weights_));
+      }
+      results.at(k) = q;
+    });
+
+    // fixme: verify this is correct for external memory
+    if (quantiles.empty()) {
+      quantiles.resize(results.size(), 0);
+    }
+    for (size_t i = 0; i < results.size(); ++i) {
+      quantiles[i] += results[i];
+    }
+  }
+
+  // use the mean value
+  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
+  auto world = rabit::GetWorldSize();
+  std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
+                 [&](float q) { return q / world; });
+
+  // fixme: verify this is correct for external memory
+  for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
+    auto seg = row_index.front().indptr[i];
+    auto q = quantiles[i];
+    CHECK(tree[seg.nidx].IsLeaf());
+    tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
+  }
+}
+
+class MeanAbsoluteError : public ObjFunction {
+ public:
+  void Configure(Args const&) override {}
+
+  uint32_t Targets(MetaInfo const& info) const override {
+    return std::max(static_cast<size_t>(1), info.labels.Shape(1));
+  }
+
+  struct ObjInfo Task() const override {
+    return {ObjInfo::kRegression, true, true};
+  }
+
+  void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
+                   HostDeviceVector<GradientPair>* out_gpair) override {
+    auto labels = info.labels.View(ctx_->gpu_id);
+
+    out_gpair->SetDevice(ctx_->gpu_id);
+    out_gpair->Resize(info.labels.Size());
+    auto gpair = linalg::MakeVec(out_gpair);
+
+    preds.SetDevice(ctx_->gpu_id);
+    auto predt = linalg::MakeVec(&preds);
+    info.weights_.SetDevice(ctx_->gpu_id);
+    common::OptionalWeights weight{ctx_->IsCPU() ? info.weights_.ConstHostSpan()
+                                                 : info.weights_.ConstDeviceSpan()};
+
+    linalg::ElementWiseKernel(ctx_, labels, [=] XGBOOST_DEVICE(size_t i, float const y) mutable {
+      auto sign = [](auto x) {
+        return (x > static_cast<decltype(x)>(0)) - (x < static_cast<decltype(x)>(0));
+      };
+      auto sample_id = std::get<0>(linalg::UnravelIndex(i, labels.Shape()));
+      auto grad = sign(predt(i) - y) * weight[i];
+      auto hess = weight[sample_id];
+      gpair(i) = GradientPair{grad, hess};
+    });
+  }
+
+  void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
+                      HostDeviceVector<float> const& prediction, uint32_t target,
+                      RegTree* p_tree) const override {
+    UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
+  }
+
+  const char* DefaultEvalMetric() const override { return "mae"; }
+
+  void SaveConfig(Json* p_out) const override {
+    auto& out = *p_out;
+    out["name"] = String("reg:absoluteerror");
+  }
+
+  void LoadConfig(Json const& in) override {}
+};
+
+XGBOOST_REGISTER_OBJECTIVE(MeanAbsoluteError, "reg:absoluteerror")
+    .describe("Mean absoluate error.")
+    .set_body([]() { return new MeanAbsoluteError(); });
 }  // namespace obj
 }  // namespace xgboost
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index cb7dd9b7e8e4..ffed619f10f5 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -6,11 +6,14 @@
 #include <xgboost/tree_updater.h>
 #include <algorithm>
 #include <cmath>
+#include <cstddef>
 #include <memory>
 #include <limits>
 #include <utility>
 #include <vector>
 
+#include "xgboost/base.h"
+#include "xgboost/generic_parameters.h"
 #include "xgboost/host_device_vector.h"
 #include "xgboost/parameter.h"
 #include "xgboost/span.h"
@@ -35,6 +38,8 @@
 #include "gpu_hist/histogram.cuh"
 #include "gpu_hist/evaluate_splits.cuh"
 #include "gpu_hist/expand_entry.cuh"
+#include "xgboost/task.h"
+#include "xgboost/tree_model.h"
 
 namespace xgboost {
 namespace tree {
@@ -384,7 +389,7 @@ struct GPUHistMakerDevice {
   // After tree update is finished, update the position of all training
   // instances to their final leaf. This information is used later to update the
   // prediction cache
-  void FinalisePosition(RegTree const* p_tree, DMatrix* p_fmat) {
+  void FinalisePosition(RegTree const* p_tree, DMatrix* p_fmat, dh::device_vector<RowIndexCache>* p_out_row_indices) {
     dh::TemporaryArray<RegTree::Node> d_nodes(p_tree->GetNodes().size());
     dh::safe_cuda(cudaMemcpyAsync(d_nodes.data().get(), p_tree->GetNodes().data(),
                                   d_nodes.size() * sizeof(RegTree::Node),
@@ -418,6 +423,12 @@ struct GPUHistMakerDevice {
                                dh::ToSpan(d_categories_segments));
       }
     }
+
+    CHECK(p_out_row_indices->empty());
+    p_out_row_indices->emplace_back(ctx_, p_fmat->Info().num_row_);
+    auto d_row_index = p_out_row_indices->back().DeviceSpan();
+    thrust::copy_if();  // compact sample nodes.
+    // run length encode.
   }
 
   void FinalisePositionInPage(EllpackPageImpl const *page,
@@ -461,7 +472,7 @@ struct GPUHistMakerDevice {
         });
   }
 
-  void UpdatePredictionCache(linalg::VectorView<float> out_preds_d) {
+  void UpdatePredictionCache(linalg::VectorView<float> out_preds_d, RegTree const* p_tree) {
     dh::safe_cuda(cudaSetDevice(device_id));
     CHECK_EQ(out_preds_d.DeviceIdx(), device_id);
     auto d_ridx = row_partitioner->GetRows();
@@ -476,13 +487,29 @@ struct GPUHistMakerDevice {
     auto d_node_sum_gradients = device_node_sum_gradients.data().get();
     auto tree_evaluator = evaluator_.GetEvaluator();
 
-    dh::LaunchN(d_ridx.size(), [=, out_preds_d = out_preds_d] __device__(int local_idx) mutable {
-      int pos = d_position[local_idx];
-      bst_float weight =
-          tree_evaluator.CalcWeight(pos, param_d, GradStats{d_node_sum_gradients[pos]});
-      static_assert(!std::is_const<decltype(out_preds_d)>::value, "");
-      out_preds_d(d_ridx[local_idx]) += weight * param_d.learning_rate;
-    });
+    if (p_tree) {
+      auto const& h_nodes = p_tree->GetNodes();
+      dh::device_vector<RegTree::Node> nodes(h_nodes.size());
+      dh::safe_cuda(cudaMemcpyAsync(nodes.data().get(), h_nodes.data(),
+                                    h_nodes.size() * sizeof(RegTree::Node),
+                                    cudaMemcpyHostToDevice));
+      auto d_nodes = dh::ToSpan(nodes);
+      dh::LaunchN(d_ridx.size(), [=] XGBOOST_DEVICE(size_t idx) mutable {
+        bst_node_t nidx = d_position[idx];
+        auto weight = d_nodes[nidx].LeafValue();
+        out_preds_d(d_ridx[idx]) += weight;
+      });
+    } else {
+      // Avoid copying nodes by using evaluator to get leaf weight on-the-fly, this is
+      // useful when tree leaf is not updated after tree construction.
+      dh::LaunchN(d_ridx.size(), [=] XGBOOST_DEVICE(size_t local_idx) mutable {
+        bst_node_t nidx = d_position[local_idx];
+        float weight =
+            tree_evaluator.CalcWeight(nidx, param_d, GradStats{d_node_sum_gradients[nidx]});
+        static_assert(!std::is_const<decltype(out_preds_d)>::value, "");
+        out_preds_d(d_ridx[local_idx]) += weight * param_d.learning_rate;
+      });
+    }
     row_partitioner.reset();
   }
 
@@ -610,7 +637,8 @@ struct GPUHistMakerDevice {
   }
 
   void UpdateTree(HostDeviceVector<GradientPair>* gpair_all, DMatrix* p_fmat, ObjInfo task,
-                  RegTree* p_tree, dh::AllReducer* reducer) {
+                  RegTree* p_tree, dh::AllReducer* reducer,
+                  dh::device_vector<RowIndexCache>* p_out_row_indices) {
     auto& tree = *p_tree;
     Driver<GPUExpandEntry> driver(static_cast<TrainParam::TreeGrowPolicy>(param.grow_policy));
 
@@ -671,7 +699,7 @@ struct GPUHistMakerDevice {
     }
 
     monitor.Start("FinalisePosition");
-    this->FinalisePosition(p_tree, p_fmat);
+    this->FinalisePosition(p_tree, p_fmat, p_out_row_indices);
     monitor.Stop("FinalisePosition");
   }
 };
@@ -682,7 +710,7 @@ class GPUHistMakerSpecialised {
   explicit GPUHistMakerSpecialised(ObjInfo task) : task_{task} {};
   void Configure(const Args& args, GenericParameter const* generic_param) {
     param_.UpdateAllowUnknown(args);
-    generic_param_ = generic_param;
+    ctx_ = generic_param;
     hist_maker_param_.UpdateAllowUnknown(args);
     dh::CheckComputeCapability();
 
@@ -703,7 +731,9 @@ class GPUHistMakerSpecialised {
     // build tree
     try {
       for (xgboost::RegTree* tree : trees) {
-        this->UpdateTree(gpair, dmat, tree);
+        row_set_collection_.emplace_back();
+        auto &row_indices = row_set_collection_.back();
+        this->UpdateTree(gpair, dmat, tree, &row_indices);
 
         if (hist_maker_param_.debug_synchronize) {
           this->CheckTreesSynchronized(tree);
@@ -719,23 +749,22 @@ class GPUHistMakerSpecialised {
   }
 
   void InitDataOnce(DMatrix* dmat) {
-    device_ = generic_param_->gpu_id;
-    CHECK_GE(device_, 0) << "Must have at least one device";
+    CHECK_GE(ctx_->gpu_id, 0) << "Must have at least one device";
     info_ = &dmat->Info();
-    reducer_.Init({device_});  // NOLINT
+    reducer_.Init({ctx_->gpu_id});  // NOLINT
 
     // Synchronise the column sampling seed
     uint32_t column_sampling_seed = common::GlobalRandom()();
     rabit::Broadcast(&column_sampling_seed, sizeof(column_sampling_seed), 0);
 
     BatchParam batch_param{
-      device_,
+      ctx_->gpu_id,
       param_.max_bin,
     };
     auto page = (*dmat->GetBatches<EllpackPage>(batch_param).begin()).Impl();
-    dh::safe_cuda(cudaSetDevice(device_));
-    info_->feature_types.SetDevice(device_);
-    maker.reset(new GPUHistMakerDevice<GradientSumT>(device_,
+    dh::safe_cuda(cudaSetDevice(ctx_->gpu_id));
+    info_->feature_types.SetDevice(ctx_->gpu_id);
+    maker.reset(new GPUHistMakerDevice<GradientSumT>(ctx_->gpu_id,
                                                      page,
                                                      info_->feature_types.ConstDeviceSpan(),
                                                      info_->num_row_,
@@ -748,12 +777,13 @@ class GPUHistMakerSpecialised {
     initialised_ = true;
   }
 
-  void InitData(DMatrix* dmat) {
+  void InitData(DMatrix* dmat, RegTree const* p_tree) {
     if (!initialised_) {
       monitor_.Start("InitDataOnce");
       this->InitDataOnce(dmat);
       monitor_.Stop("InitDataOnce");
     }
+    p_last_tree_ = p_tree;
   }
 
   // Only call this method for testing
@@ -771,12 +801,13 @@ class GPUHistMakerSpecialised {
     CHECK(*local_tree == reference_tree);
   }
 
-  void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat, RegTree* p_tree) {
+  void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat, RegTree* p_tree,
+                  dh::device_vector<RowIndexCache>* p_out_row_indices) {
     monitor_.Start("InitData");
-    this->InitData(p_fmat);
+    this->InitData(p_fmat, p_tree);
     monitor_.Stop("InitData");
 
-    gpair->SetDevice(device_);
+    gpair->SetDevice(ctx_->gpu_id);
     maker->UpdateTree(gpair, p_fmat, task_, p_tree, &reducer_);
   }
 
@@ -786,7 +817,7 @@ class GPUHistMakerSpecialised {
       return false;
     }
     monitor_.Start("UpdatePredictionCache");
-    maker->UpdatePredictionCache(p_out_preds);
+    maker->UpdatePredictionCache(p_out_preds, task_.zero_hess ? p_last_tree_ : nullptr);
     monitor_.Stop("UpdatePredictionCache");
     return true;
   }
@@ -798,14 +829,15 @@ class GPUHistMakerSpecialised {
 
  private:
   bool initialised_ { false };
+  std::vector<dh::device_vector<RowIndexCache>> row_set_collection_;
 
   GPUHistMakerTrainParam hist_maker_param_;
-  GenericParameter const* generic_param_;
+  Context const* ctx_;
 
   dh::AllReducer reducer_;
 
   DMatrix* p_last_fmat_ { nullptr };
-  int device_{-1};
+  RegTree const* p_last_tree_{nullptr};
   ObjInfo task_;
 
   common::Monitor monitor_;
@@ -867,9 +899,8 @@ class GPUHistMaker : public TreeUpdater {
     }
   }
 
-  bool
-  UpdatePredictionCache(const DMatrix *data,
-                        linalg::VectorView<bst_float> p_out_preds) override {
+  bool UpdatePredictionCache(const DMatrix* data,
+                             linalg::VectorView<bst_float> p_out_preds) override {
     if (hist_maker_param_.single_precision_histogram) {
       return float_maker_->UpdatePredictionCache(data, p_out_preds);
     } else {
@@ -877,6 +908,10 @@ class GPUHistMaker : public TreeUpdater {
     }
   }
 
+  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
+    return dh::ToSpan(row_set_collection_.at(tree_idx));
+  }
+
   char const* Name() const override {
     return "grow_gpu_hist";
   }
@@ -884,6 +919,7 @@ class GPUHistMaker : public TreeUpdater {
  private:
   GPUHistMakerTrainParam hist_maker_param_;
   ObjInfo task_;
+  std::vector<dh::device_vector<RowIndexCache>> row_set_collection_;
   std::unique_ptr<GPUHistMakerSpecialised<GradientPair>> float_maker_;
   std::unique_ptr<GPUHistMakerSpecialised<GradientPairPrecise>> double_maker_;
 };

From 9513cf446e3a67a55189719aba24bce343bd5519 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Sat, 9 Apr 2022 11:54:52 +0800
Subject: [PATCH 018/124] Copy the row index.

---
 src/tree/gpu_hist/row_partitioner.cuh |  2 +-
 src/tree/updater_gpu_hist.cu          | 67 ++++++++++++++++-----------
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 1b5a5222229e..5ff3c211f39c 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -34,6 +34,7 @@ class RowPartitioner {
   using RowIndexT = bst_uint;
   struct Segment;
   static constexpr bst_node_t kIgnoredTreePosition = -1;
+  std::vector<Segment> ridx_segments_;
 
  private:
   int device_idx_;
@@ -45,7 +46,6 @@ class RowPartitioner {
    * node id -> segment -> indices of rows belonging to node
    */
   /*! \brief Range of row index for each node, pointers into ridx below. */
-  std::vector<Segment> ridx_segments_;
   dh::TemporaryArray<RowIndexT> ridx_a_;
   dh::TemporaryArray<RowIndexT> ridx_b_;
   dh::TemporaryArray<bst_node_t> position_a_;
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index ffed619f10f5..fba3d55d3cfc 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -166,9 +166,9 @@ template <typename GradientSumT>
 struct GPUHistMakerDevice {
  private:
   GPUHistEvaluator<GradientSumT> evaluator_;
+  Context const* ctx_;
 
  public:
-  int device_id;
   EllpackPageImpl const* page;
   common::Span<FeatureType const> feature_types;
   BatchParam batch_param;
@@ -200,12 +200,12 @@ struct GPUHistMakerDevice {
   // Storing split categories for last node.
   dh::caching_device_vector<uint32_t> node_categories;
 
-  GPUHistMakerDevice(int _device_id, EllpackPageImpl const* _page,
+  GPUHistMakerDevice(Context const* ctx, EllpackPageImpl const* _page,
                      common::Span<FeatureType const> _feature_types, bst_uint _n_rows,
                      TrainParam _param, uint32_t column_sampler_seed, uint32_t n_features,
                      BatchParam _batch_param)
-      : evaluator_{_param, n_features, _device_id},
-        device_id(_device_id),
+      : evaluator_{_param, n_features, ctx->gpu_id},
+        ctx_(ctx),
         page(_page),
         feature_types{_feature_types},
         param(std::move(_param)),
@@ -221,14 +221,14 @@ struct GPUHistMakerDevice {
     node_sum_gradients.resize(param.MaxNodes());
 
     // Init histogram
-    hist.Init(device_id, page->Cuts().TotalBins());
-    monitor.Init(std::string("GPUHistMakerDevice") + std::to_string(device_id));
+    hist.Init(ctx_->gpu_id, page->Cuts().TotalBins());
+    monitor.Init(std::string("GPUHistMakerDevice") + std::to_string(ctx_->gpu_id));
     feature_groups.reset(new FeatureGroups(
-        page->Cuts(), page->is_dense, dh::MaxSharedMemoryOptin(device_id), sizeof(GradientSumT)));
+        page->Cuts(), page->is_dense, dh::MaxSharedMemoryOptin(ctx_->gpu_id), sizeof(GradientSumT)));
   }
 
   ~GPUHistMakerDevice() {  // NOLINT
-    dh::safe_cuda(cudaSetDevice(device_id));
+    dh::safe_cuda(cudaSetDevice(ctx_->gpu_id));
   }
 
   // Reset values for each update iteration
@@ -240,10 +240,10 @@ struct GPUHistMakerDevice {
     this->column_sampler.Init(num_columns, info.feature_weights.HostVector(),
                               param.colsample_bynode, param.colsample_bylevel,
                               param.colsample_bytree);
-    dh::safe_cuda(cudaSetDevice(device_id));
+    dh::safe_cuda(cudaSetDevice(ctx_->gpu_id));
 
     this->evaluator_.Reset(page->Cuts(), feature_types, task, dmat->Info().num_col_, param,
-                           device_id);
+                           ctx_->gpu_id);
 
     this->interaction_constraints.Reset();
     std::fill(node_sum_gradients.begin(), node_sum_gradients.end(), GradientPairPrecise{});
@@ -261,7 +261,7 @@ struct GPUHistMakerDevice {
     histogram_rounding = CreateRoundingFactor<GradientSumT>(this->gpair);
 
     row_partitioner.reset();  // Release the device memory first before reallocating
-    row_partitioner.reset(new RowPartitioner(device_id,  sample.sample_rows));
+    row_partitioner.reset(new RowPartitioner(ctx_->gpu_id,  sample.sample_rows));
     hist.Reset();
   }
 
@@ -269,10 +269,10 @@ struct GPUHistMakerDevice {
     int nidx = RegTree::kRoot;
     GPUTrainingParam gpu_param(param);
     auto sampled_features = column_sampler.GetFeatureSet(0);
-    sampled_features->SetDevice(device_id);
+    sampled_features->SetDevice(ctx_->gpu_id);
     common::Span<bst_feature_t> feature_set =
         interaction_constraints.Query(sampled_features->DeviceSpan(), nidx);
-    auto matrix = page->GetDeviceAccessor(device_id);
+    auto matrix = page->GetDeviceAccessor(ctx_->gpu_id);
     EvaluateSplitInputs<GradientSumT> inputs{nidx,
                                              root_sum,
                                              gpu_param,
@@ -292,14 +292,14 @@ struct GPUHistMakerDevice {
     dh::TemporaryArray<DeviceSplitCandidate> splits_out(2);
     GPUTrainingParam gpu_param(param);
     auto left_sampled_features = column_sampler.GetFeatureSet(tree.GetDepth(left_nidx));
-    left_sampled_features->SetDevice(device_id);
+    left_sampled_features->SetDevice(ctx_->gpu_id);
     common::Span<bst_feature_t> left_feature_set =
         interaction_constraints.Query(left_sampled_features->DeviceSpan(), left_nidx);
     auto right_sampled_features = column_sampler.GetFeatureSet(tree.GetDepth(right_nidx));
-    right_sampled_features->SetDevice(device_id);
+    right_sampled_features->SetDevice(ctx_->gpu_id);
     common::Span<bst_feature_t> right_feature_set =
         interaction_constraints.Query(right_sampled_features->DeviceSpan(), left_nidx);
-    auto matrix = page->GetDeviceAccessor(device_id);
+    auto matrix = page->GetDeviceAccessor(ctx_->gpu_id);
 
     EvaluateSplitInputs<GradientSumT> left{left_nidx,
                                            candidate.split.left_sum,
@@ -330,8 +330,8 @@ struct GPUHistMakerDevice {
     hist.AllocateHistogram(nidx);
     auto d_node_hist = hist.GetNodeHistogram(nidx);
     auto d_ridx = row_partitioner->GetRows(nidx);
-    BuildGradientHistogram(page->GetDeviceAccessor(device_id),
-                           feature_groups->DeviceAccessor(device_id), gpair,
+    BuildGradientHistogram(page->GetDeviceAccessor(ctx_->gpu_id),
+                           feature_groups->DeviceAccessor(ctx_->gpu_id), gpair,
                            d_ridx, d_node_hist, histogram_rounding);
   }
 
@@ -356,7 +356,7 @@ struct GPUHistMakerDevice {
   void UpdatePosition(int nidx, RegTree* p_tree) {
     RegTree::Node split_node = (*p_tree)[nidx];
     auto split_type = p_tree->NodeSplitType(nidx);
-    auto d_matrix = page->GetDeviceAccessor(device_id);
+    auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
     auto node_cats = dh::ToSpan(node_categories);
 
     row_partitioner->UpdatePosition(
@@ -389,7 +389,7 @@ struct GPUHistMakerDevice {
   // After tree update is finished, update the position of all training
   // instances to their final leaf. This information is used later to update the
   // prediction cache
-  void FinalisePosition(RegTree const* p_tree, DMatrix* p_fmat, dh::device_vector<RowIndexCache>* p_out_row_indices) {
+  void FinalisePosition(RegTree const* p_tree, DMatrix* p_fmat, std::vector<RowIndexCache>* p_out_row_indices) {
     dh::TemporaryArray<RegTree::Node> d_nodes(p_tree->GetNodes().size());
     dh::safe_cuda(cudaMemcpyAsync(d_nodes.data().get(), p_tree->GetNodes().data(),
                                   d_nodes.size() * sizeof(RegTree::Node),
@@ -410,7 +410,7 @@ struct GPUHistMakerDevice {
 
     if (row_partitioner->GetRows().size() != p_fmat->Info().num_row_) {
       row_partitioner.reset();  // Release the device memory first before reallocating
-      row_partitioner.reset(new RowPartitioner(device_id, p_fmat->Info().num_row_));
+      row_partitioner.reset(new RowPartitioner(ctx_->gpu_id, p_fmat->Info().num_row_));
     }
     if (page->n_rows == p_fmat->Info().num_row_) {
       FinalisePositionInPage(page, dh::ToSpan(d_nodes),
@@ -425,10 +425,21 @@ struct GPUHistMakerDevice {
     }
 
     CHECK(p_out_row_indices->empty());
-    p_out_row_indices->emplace_back(ctx_, p_fmat->Info().num_row_);
-    auto d_row_index = p_out_row_indices->back().DeviceSpan();
-    thrust::copy_if();  // compact sample nodes.
-    // run length encode.
+    p_out_row_indices->push_back(RowIndexCache{ctx_, p_fmat->Info().num_row_});
+    auto& segments = p_out_row_indices->back().indptr;
+    auto const& tree = *p_tree;
+    for (size_t nidx = 0; nidx < row_partitioner->ridx_segments_.size(); ++nidx) {
+      if (tree[nidx].IsLeaf()) {
+        auto const& seg = row_partitioner->ridx_segments_[nidx];
+        // fixme: subsample
+        segments.emplace_back(seg.begin, seg.Size(), static_cast<bst_node_t>(nidx));
+      }
+    }
+    auto in = row_partitioner->GetRows();
+    p_out_row_indices->back().row_index.Resize(in.size());
+    auto d_row_index = p_out_row_indices->back().row_index.DeviceSpan();
+    dh::safe_cuda(
+        cudaMemcpyAsync(d_row_index.data(), in.data(), in.size_bytes(), cudaMemcpyDeviceToDevice));
   }
 
   void FinalisePositionInPage(EllpackPageImpl const *page,
@@ -436,7 +447,7 @@ struct GPUHistMakerDevice {
                               common::Span<FeatureType const> d_feature_types,
                               common::Span<uint32_t const> categories,
                               common::Span<RegTree::Segment> categories_segments) {
-    auto d_matrix = page->GetDeviceAccessor(device_id);
+    auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
     row_partitioner->FinalisePosition(
         [=] __device__(size_t row_id, int position) {
           // What happens if user prune the tree?
@@ -473,8 +484,8 @@ struct GPUHistMakerDevice {
   }
 
   void UpdatePredictionCache(linalg::VectorView<float> out_preds_d, RegTree const* p_tree) {
-    dh::safe_cuda(cudaSetDevice(device_id));
-    CHECK_EQ(out_preds_d.DeviceIdx(), device_id);
+    dh::safe_cuda(cudaSetDevice(ctx_->gpu_id));
+    CHECK_EQ(out_preds_d.DeviceIdx(), ctx_->gpu_id);
     auto d_ridx = row_partitioner->GetRows();
 
     GPUTrainingParam param_d(param);

From d9ac306351bd69b24522f8b842c315b4c70fac92 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Sat, 9 Apr 2022 12:14:28 +0800
Subject: [PATCH 019/124] Comp.

---
 src/objective/regression_obj.cu |  8 +++++++-
 src/tree/updater_gpu_hist.cu    | 21 ++++++++-------------
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 324a8d166b74..ad9114b05860 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -678,6 +678,8 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   dh::LaunchN(part.row_index.Size(), [=]XGBOOST_DEVICE(size_t i) {
 
   });
+
+  LOG(FATAL) << "Not implemented";
 }
 
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
@@ -773,7 +775,11 @@ class MeanAbsoluteError : public ObjFunction {
   void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
                       HostDeviceVector<float> const& prediction, uint32_t target,
                       RegTree* p_tree) const override {
-    UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
+    if (ctx_->IsCPU()) {
+      UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
+    } else {
+      UpdateTreeLeafDevice(ctx_, row_index, info, prediction, target, 0.5, p_tree);
+    }
   }
 
   const char* DefaultEvalMetric() const override { return "mae"; }
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index fba3d55d3cfc..3f5f4ebd74fc 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -649,7 +649,7 @@ struct GPUHistMakerDevice {
 
   void UpdateTree(HostDeviceVector<GradientPair>* gpair_all, DMatrix* p_fmat, ObjInfo task,
                   RegTree* p_tree, dh::AllReducer* reducer,
-                  dh::device_vector<RowIndexCache>* p_out_row_indices) {
+                  std::vector<RowIndexCache>* p_out_row_indices) {
     auto& tree = *p_tree;
     Driver<GPUExpandEntry> driver(static_cast<TrainParam::TreeGrowPolicy>(param.grow_policy));
 
@@ -743,7 +743,7 @@ class GPUHistMakerSpecialised {
     try {
       for (xgboost::RegTree* tree : trees) {
         row_set_collection_.emplace_back();
-        auto &row_indices = row_set_collection_.back();
+        auto& row_indices = row_set_collection_.back();
         this->UpdateTree(gpair, dmat, tree, &row_indices);
 
         if (hist_maker_param_.debug_synchronize) {
@@ -775,14 +775,9 @@ class GPUHistMakerSpecialised {
     auto page = (*dmat->GetBatches<EllpackPage>(batch_param).begin()).Impl();
     dh::safe_cuda(cudaSetDevice(ctx_->gpu_id));
     info_->feature_types.SetDevice(ctx_->gpu_id);
-    maker.reset(new GPUHistMakerDevice<GradientSumT>(ctx_->gpu_id,
-                                                     page,
-                                                     info_->feature_types.ConstDeviceSpan(),
-                                                     info_->num_row_,
-                                                     param_,
-                                                     column_sampling_seed,
-                                                     info_->num_col_,
-                                                     batch_param));
+    maker.reset(new GPUHistMakerDevice<GradientSumT>(
+        ctx_, page, info_->feature_types.ConstDeviceSpan(), info_->num_row_, param_,
+        column_sampling_seed, info_->num_col_, batch_param));
 
     p_last_fmat_ = dmat;
     initialised_ = true;
@@ -813,13 +808,13 @@ class GPUHistMakerSpecialised {
   }
 
   void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat, RegTree* p_tree,
-                  dh::device_vector<RowIndexCache>* p_out_row_indices) {
+                  std::vector<RowIndexCache>* p_out_row_indices) {
     monitor_.Start("InitData");
     this->InitData(p_fmat, p_tree);
     monitor_.Stop("InitData");
 
     gpair->SetDevice(ctx_->gpu_id);
-    maker->UpdateTree(gpair, p_fmat, task_, p_tree, &reducer_);
+    maker->UpdateTree(gpair, p_fmat, task_, p_tree, &reducer_, p_out_row_indices);
   }
 
   bool UpdatePredictionCache(const DMatrix *data,
@@ -840,7 +835,7 @@ class GPUHistMakerSpecialised {
 
  private:
   bool initialised_ { false };
-  std::vector<dh::device_vector<RowIndexCache>> row_set_collection_;
+  std::vector<std::vector<RowIndexCache>> row_set_collection_;
 
   GPUHistMakerTrainParam hist_maker_param_;
   Context const* ctx_;

From eca5afbbf1b23f920db6e313cc328f547c4d0210 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Sat, 9 Apr 2022 12:57:00 +0800
Subject: [PATCH 020/124] Start working on device quantile.

---
 src/objective/regression_obj.cu | 74 +++++++++++++++++++++++++++++----
 1 file changed, 66 insertions(+), 8 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index ad9114b05860..95791eb6d1c9 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -11,6 +11,7 @@
 #include <xgboost/tree_model.h>
 
 #include <cmath>
+#include <cstddef>
 #include <memory>
 #include <vector>
 
@@ -21,6 +22,8 @@
 #include "../common/threading_utils.h"
 #include "../common/transform.h"
 #include "./regression_loss.h"
+#include "xgboost/base.h"
+#include "xgboost/data.h"
 #include "xgboost/generic_parameters.h"
 #include "xgboost/host_device_vector.h"
 #include "xgboost/json.h"
@@ -664,6 +667,59 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .describe("Tweedie regression for insurance data.")
 .set_body([]() { return new TweedieRegression(); });
 
+void SegmentedPercentile(double alpha, RowIndexCache const& row_index, MetaInfo const& info,
+                         HostDeviceVector<float> const& predt, HostDeviceVector<float>* quantiles) {
+  CHECK(alpha >= 0 && alpha <= 1);
+  dh::device_vector<float> residue(predt.Size());
+  auto d_residue = dh::ToSpan(residue);
+  auto d_predt = predt.ConstDeviceSpan();
+  auto d_labels = info.labels.View(0);
+  dh::LaunchN(residue.size(), [=] XGBOOST_DEVICE(size_t i) {
+    // compute residule
+  });
+
+  dh::XGBDeviceAllocator<char> alloc;
+  dh::device_vector<size_t> sorted_idx(residue.size());
+  dh::Iota(dh::ToSpan(sorted_idx));
+  using Tup = thrust::tuple<size_t, float>;
+  auto key_it = dh::MakeTransformIterator<size_t>(thrust::make_counting_iterator(0ul),
+                                                  [=] XGBOOST_DEVICE(size_t i) -> Tup {});
+  thrust::stable_sort_by_key(thrust::cuda::par(alloc), key_it, key_it + residue.size(),
+                             sorted_idx.begin(), [=] XGBOOST_DEVICE(Tup const& l, Tup const& r) {
+                               if (thrust::get<0>(l) != thrust::get<0>(r)) {
+                                 return thrust::get<0>(l) < thrust::get<0>(r);  // segment index
+                               }
+                               return thrust::get<1>(l) < thrust::get<1>(r);  // residue
+                             });
+
+  dh::caching_device_vector<RowIndexCache::Segment> segment(row_index.indptr.size());
+  thrust::copy(row_index.indptr.cbegin(), row_index.indptr.cend(), segment.begin());
+  auto d_segments = dh::ToSpan(segment);
+
+  quantiles->Resize(row_index.indptr.size());
+  auto d_results = quantiles->DeviceSpan();
+
+  auto d_row_index = row_index.row_index.ConstDeviceSpan();
+  auto d_sorted_idx = dh::ToSpan(sorted_idx);
+
+  dh::LaunchN(residue.size(), [=] XGBOOST_DEVICE(size_t i) {
+    // each segment is the index of a leaf.
+    size_t seg_idx = 0;
+    auto seg = d_segments[seg_idx];
+    auto residue_seg = d_residue.subspan(seg.begin, seg.n);
+
+    double x = alpha * static_cast<double>(seg.n + 1);
+    double k = std::floor(x) - 1;
+    double d = (x - 1) - k;
+
+    if (i == seg.begin) {
+      auto v0 = d_residue[d_row_index[d_sorted_idx[static_cast<size_t>(k)]]];
+      auto v1 = d_residue[d_row_index[d_sorted_idx[static_cast<size_t>(k) + 1]]];
+      d_results[seg_idx] = v0 + d * (v1 - v0);
+    }
+  });
+}
+
 void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const> row_index,
                           MetaInfo const& info, HostDeviceVector<float> const& prediction,
                           uint32_t target, float alpha, RegTree* p_tree) {
@@ -671,15 +727,17 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   CHECK_EQ(row_index.size(), 1);
   auto const& part = row_index.front();
 
-  dh::caching_device_vector<float> quantiles;
-  auto d_quantiles = dh::ToSpan(quantiles);
+  HostDeviceVector<float> results;
+  SegmentedPercentile(alpha, part, info, prediction, &results);
 
-  auto rows = part.row_index.ConstDeviceSpan();
-  dh::LaunchN(part.row_index.Size(), [=]XGBOOST_DEVICE(size_t i) {
-
-  });
-
-  LOG(FATAL) << "Not implemented";
+  auto const& h_results = results.HostVector();
+  auto& tree = *p_tree;
+  for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
+    auto seg = row_index.front().indptr[i];
+    auto q = h_results[i];
+    CHECK(tree[seg.nidx].IsLeaf());
+    tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
+  }
 }
 
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,

From 848088617019c76849cb6085b5b2a730736cd453 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Tue, 12 Apr 2022 13:56:27 +0800
Subject: [PATCH 021/124] Compute target.

---
 include/xgboost/linalg.h        |  9 ++++-
 src/common/linalg_op.cuh        |  2 +-
 src/objective/regression_obj.cu | 69 ++++++++++++++++++++++++---------
 src/tree/updater_gpu_hist.cu    |  7 ++--
 4 files changed, 62 insertions(+), 25 deletions(-)

diff --git a/include/xgboost/linalg.h b/include/xgboost/linalg.h
index 32d0f9fb9f9c..15b3a04235bf 100644
--- a/include/xgboost/linalg.h
+++ b/include/xgboost/linalg.h
@@ -14,6 +14,7 @@
 
 #include <algorithm>
 #include <cassert>
+#include <cstddef>
 #include <limits>
 #include <string>
 #include <type_traits>
@@ -670,9 +671,13 @@ class Tensor {
    * See \ref TensorView for parameters of this constructor.
    */
   template <typename I, int32_t D>
-  explicit Tensor(I const (&shape)[D], int32_t device) {
+  explicit Tensor(I const (&shape)[D], int32_t device)
+      : Tensor{common::Span<I, D>{shape}, device} {}
+
+  template <typename I, size_t D>
+  explicit Tensor(common::Span<I, D> shape, int32_t device) {
     // No device unroll as this is a host only function.
-    std::copy(shape, shape + D, shape_);
+    std::copy(shape.data(), shape.data() + D, shape_);
     for (auto i = D; i < kDim; ++i) {
       shape_[i] = 1;
     }
diff --git a/src/common/linalg_op.cuh b/src/common/linalg_op.cuh
index f0f89df8ab26..84caae2670cd 100644
--- a/src/common/linalg_op.cuh
+++ b/src/common/linalg_op.cuh
@@ -40,7 +40,7 @@ void ElementWiseTransformDevice(linalg::TensorView<T, D> t, Fn&& fn, cudaStream_
 }
 
 template <typename T, int32_t D, typename Fn>
-void ElementWiseKernel(GenericParameter const* ctx, linalg::TensorView<T, D> t, Fn&& fn) {
+void ElementWiseKernel(Context const* ctx, linalg::TensorView<T, D> t, Fn&& fn) {
   ctx->IsCPU() ? ElementWiseKernelHost(t, ctx->Threads(), fn) : ElementWiseKernelDevice(t, fn);
 }
 }  // namespace linalg
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 95791eb6d1c9..ced9a450b4f8 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -27,6 +27,7 @@
 #include "xgboost/generic_parameters.h"
 #include "xgboost/host_device_vector.h"
 #include "xgboost/json.h"
+#include "xgboost/linalg.h"
 #include "xgboost/parameter.h"
 #include "xgboost/span.h"
 
@@ -667,25 +668,52 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .describe("Tweedie regression for insurance data.")
 .set_body([]() { return new TweedieRegression(); });
 
-void SegmentedPercentile(double alpha, RowIndexCache const& row_index, MetaInfo const& info,
-                         HostDeviceVector<float> const& predt, HostDeviceVector<float>* quantiles) {
+void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const& row_index,
+                         MetaInfo const& info, HostDeviceVector<float> const& predt,
+                         HostDeviceVector<float>* quantiles) {
   CHECK(alpha >= 0 && alpha <= 1);
-  dh::device_vector<float> residue(predt.Size());
-  auto d_residue = dh::ToSpan(residue);
+
   auto d_predt = predt.ConstDeviceSpan();
-  auto d_labels = info.labels.View(0);
-  dh::LaunchN(residue.size(), [=] XGBOOST_DEVICE(size_t i) {
-    // compute residule
+  auto d_labels = info.labels.View(ctx->gpu_id);
+  linalg::Tensor<float, 2> residue{d_labels.Shape(), ctx->gpu_id};
+  auto d_residue = residue.View(ctx->gpu_id);
+  CHECK_EQ(d_predt.size(), d_labels.Size());
+  linalg::ElementWiseKernel(ctx, d_labels, [=] XGBOOST_DEVICE(size_t i, float y) mutable {
+    size_t sample_id, target_id;
+    std::tie(sample_id, target_id) = linalg::UnravelIndex(i, d_labels.Shape());
+    d_residue(sample_id, target_id) = y - d_predt[i];
+  });
+
+  dh::device_vector<size_t> segment_idx(row_index.indptr.size() + 1, 0);
+  auto d_segment_idx = dh::ToSpan(segment_idx);
+  dh::device_vector<RowIndexCache::Segment> indptr(row_index.indptr);
+  auto d_indptr = dh::ToSpan(indptr);
+  dh::LaunchN(d_segment_idx.size(), [=] XGBOOST_DEVICE(size_t i) {
+    if (i == d_segment_idx.size() - 1) {
+      d_segment_idx[i] = d_indptr[i].begin + d_indptr[i].n;
+      return;
+    }
+    d_segment_idx[i] = d_indptr[i].begin;
   });
 
   dh::XGBDeviceAllocator<char> alloc;
-  dh::device_vector<size_t> sorted_idx(residue.size());
+  dh::device_vector<size_t> sorted_idx(d_labels.Shape(0));
   dh::Iota(dh::ToSpan(sorted_idx));
   using Tup = thrust::tuple<size_t, float>;
-  auto key_it = dh::MakeTransformIterator<size_t>(thrust::make_counting_iterator(0ul),
-                                                  [=] XGBOOST_DEVICE(size_t i) -> Tup {});
-  thrust::stable_sort_by_key(thrust::cuda::par(alloc), key_it, key_it + residue.size(),
-                             sorted_idx.begin(), [=] XGBOOST_DEVICE(Tup const& l, Tup const& r) {
+  auto key_it = dh::MakeTransformIterator<Tup>(
+      thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
+        size_t sample_id, target_id;
+        std::tie(sample_id, target_id) = linalg::UnravelIndex(i, d_labels.Shape());
+        auto leaf_idx = dh::SegmentId(d_segment_idx, sample_id);
+        auto residue = d_residue(sample_id, target_id);
+        return thrust::make_tuple(leaf_idx, residue);
+      });
+  dh::device_vector<Tup> keys(residue.Size());
+  dh::XGBCachingDeviceAllocator<char> caching;
+  thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
+
+  thrust::stable_sort_by_key(thrust::cuda::par(alloc), keys.begin(), keys.end(), sorted_idx.begin(),
+                             [=] XGBOOST_DEVICE(Tup const& l, Tup const& r) {
                                if (thrust::get<0>(l) != thrust::get<0>(r)) {
                                  return thrust::get<0>(l) < thrust::get<0>(r);  // segment index
                                }
@@ -701,20 +729,22 @@ void SegmentedPercentile(double alpha, RowIndexCache const& row_index, MetaInfo
 
   auto d_row_index = row_index.row_index.ConstDeviceSpan();
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
+  auto d_keys = dh::ToSpan(keys);
 
-  dh::LaunchN(residue.size(), [=] XGBOOST_DEVICE(size_t i) {
+  dh::LaunchN(residue.Size(), [=] XGBOOST_DEVICE(size_t i) {
+    size_t sample_id, target_id;
+    std::tie(sample_id, target_id) = linalg::UnravelIndex(i, d_labels.Shape());
     // each segment is the index of a leaf.
-    size_t seg_idx = 0;
+    size_t seg_idx = thrust::get<0>(d_keys[i]);
     auto seg = d_segments[seg_idx];
-    auto residue_seg = d_residue.subspan(seg.begin, seg.n);
 
     double x = alpha * static_cast<double>(seg.n + 1);
     double k = std::floor(x) - 1;
     double d = (x - 1) - k;
 
     if (i == seg.begin) {
-      auto v0 = d_residue[d_row_index[d_sorted_idx[static_cast<size_t>(k)]]];
-      auto v1 = d_residue[d_row_index[d_sorted_idx[static_cast<size_t>(k) + 1]]];
+      auto v0 = d_residue(d_row_index[d_sorted_idx[static_cast<size_t>(k)]], target_id);
+      auto v1 = d_residue(d_row_index[d_sorted_idx[static_cast<size_t>(k) + 1]], target_id);
       d_results[seg_idx] = v0 + d * (v1 - v0);
     }
   });
@@ -724,11 +754,12 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
                           MetaInfo const& info, HostDeviceVector<float> const& prediction,
                           uint32_t target, float alpha, RegTree* p_tree) {
   dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
-  CHECK_EQ(row_index.size(), 1);
+  CHECK_EQ(row_index.size(), 1)
+      << "External memory with GPU hist should have only 1 row partition.";
   auto const& part = row_index.front();
 
   HostDeviceVector<float> results;
-  SegmentedPercentile(alpha, part, info, prediction, &results);
+  SegmentedPercentile(ctx, alpha, part, info, prediction, &results);
 
   auto const& h_results = results.HostVector();
   auto& tree = *p_tree;
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 3f5f4ebd74fc..383c8b90c399 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -432,7 +432,8 @@ struct GPUHistMakerDevice {
       if (tree[nidx].IsLeaf()) {
         auto const& seg = row_partitioner->ridx_segments_[nidx];
         // fixme: subsample
-        segments.emplace_back(seg.begin, seg.Size(), static_cast<bst_node_t>(nidx));
+        segments.push_back(
+            RowIndexCache::Segment{seg.begin, seg.Size(), static_cast<bst_node_t>(nidx)});
       }
     }
     auto in = row_partitioner->GetRows();
@@ -915,7 +916,7 @@ class GPUHistMaker : public TreeUpdater {
   }
 
   common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
-    return dh::ToSpan(row_set_collection_.at(tree_idx));
+    return row_set_collection_.at(tree_idx);
   }
 
   char const* Name() const override {
@@ -925,7 +926,7 @@ class GPUHistMaker : public TreeUpdater {
  private:
   GPUHistMakerTrainParam hist_maker_param_;
   ObjInfo task_;
-  std::vector<dh::device_vector<RowIndexCache>> row_set_collection_;
+  std::vector<std::vector<RowIndexCache>> row_set_collection_;
   std::unique_ptr<GPUHistMakerSpecialised<GradientPair>> float_maker_;
   std::unique_ptr<GPUHistMakerSpecialised<GradientPairPrecise>> double_maker_;
 };

From 9ff0a01a4270409905840b2acf74c113ff4178e2 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Tue, 12 Apr 2022 16:59:24 +0800
Subject: [PATCH 022/124] Work on GPU partitioner.

---
 src/objective/regression_obj.cu | 54 +++++++++++++++++++++++++--------
 src/tree/updater_gpu_hist.cu    | 28 ++++++++++++++---
 2 files changed, 65 insertions(+), 17 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index ced9a450b4f8..d66bbfa26525 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -679,18 +679,32 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
   auto d_residue = residue.View(ctx->gpu_id);
   CHECK_EQ(d_predt.size(), d_labels.Size());
   linalg::ElementWiseKernel(ctx, d_labels, [=] XGBOOST_DEVICE(size_t i, float y) mutable {
-    size_t sample_id, target_id;
-    std::tie(sample_id, target_id) = linalg::UnravelIndex(i, d_labels.Shape());
+    auto idx = linalg::UnravelIndex(i, d_labels.Shape());
+    size_t sample_id = std::get<0>(idx);
+    size_t target_id = std::get<1>(idx);
     d_residue(sample_id, target_id) = y - d_predt[i];
   });
 
   dh::device_vector<size_t> segment_idx(row_index.indptr.size() + 1, 0);
   auto d_segment_idx = dh::ToSpan(segment_idx);
   dh::device_vector<RowIndexCache::Segment> indptr(row_index.indptr);
+  for (auto const& seg : row_index.indptr) {
+    std::cout << seg.nidx << ", begin:" << seg.begin  << ", n:" << seg.n << std::endl;
+  }
+  std::cout << std::endl;
+
+  std::cout << "device segment" << std::endl;
+  for (size_t i = 0; i < indptr.size(); ++i) {
+    RowIndexCache::Segment seg = indptr[i];
+    std::cout << seg.nidx << ", begin:" << seg.begin  << ", n:" << seg.n << std::endl;
+  }
+  std::cout << std::endl;
+
   auto d_indptr = dh::ToSpan(indptr);
   dh::LaunchN(d_segment_idx.size(), [=] XGBOOST_DEVICE(size_t i) {
     if (i == d_segment_idx.size() - 1) {
-      d_segment_idx[i] = d_indptr[i].begin + d_indptr[i].n;
+      d_segment_idx[i] = d_indptr[i - 1].begin + d_indptr[i - 1].n;
+      printf("last: %lu, nidx: %d, %lu, %lu, i: %lu, \n", d_segment_idx[i], d_indptr[i - 1].nidx, d_indptr[i - 1].begin, d_indptr[i - i].n, i);
       return;
     }
     d_segment_idx[i] = d_indptr[i].begin;
@@ -699,11 +713,20 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
   dh::XGBDeviceAllocator<char> alloc;
   dh::device_vector<size_t> sorted_idx(d_labels.Shape(0));
   dh::Iota(dh::ToSpan(sorted_idx));
+  {
+    std::vector<size_t> h_segments(segment_idx.size());
+    thrust::copy(segment_idx.begin(), segment_idx.end(), h_segments.begin());
+    // for (auto idx : h_segments) {
+    //   std::cout << idx << ", ";
+    // }
+    // std::cout << std::endl;
+  }
   using Tup = thrust::tuple<size_t, float>;
   auto key_it = dh::MakeTransformIterator<Tup>(
       thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
-        size_t sample_id, target_id;
-        std::tie(sample_id, target_id) = linalg::UnravelIndex(i, d_labels.Shape());
+        auto idx = linalg::UnravelIndex(i, d_labels.Shape());
+        size_t sample_id = std::get<0>(idx);
+        size_t target_id = std::get<1>(idx);
         auto leaf_idx = dh::SegmentId(d_segment_idx, sample_id);
         auto residue = d_residue(sample_id, target_id);
         return thrust::make_tuple(leaf_idx, residue);
@@ -720,29 +743,34 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
                                return thrust::get<1>(l) < thrust::get<1>(r);  // residue
                              });
 
-  dh::caching_device_vector<RowIndexCache::Segment> segment(row_index.indptr.size());
-  thrust::copy(row_index.indptr.cbegin(), row_index.indptr.cend(), segment.begin());
-  auto d_segments = dh::ToSpan(segment);
-
+  quantiles->SetDevice(ctx->gpu_id);
   quantiles->Resize(row_index.indptr.size());
   auto d_results = quantiles->DeviceSpan();
 
   auto d_row_index = row_index.row_index.ConstDeviceSpan();
+  std::cout << "Row index" << std::endl;
+  auto const& h_row_index = row_index.row_index.HostVector();
+  for (auto idx : h_row_index) {
+    std::cout << idx << ", ";
+  }
+  std::cout << std::endl;
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
   auto d_keys = dh::ToSpan(keys);
 
   dh::LaunchN(residue.Size(), [=] XGBOOST_DEVICE(size_t i) {
-    size_t sample_id, target_id;
-    std::tie(sample_id, target_id) = linalg::UnravelIndex(i, d_labels.Shape());
+    auto idx = linalg::UnravelIndex(i, d_labels.Shape());
+    size_t target_id = std::get<1>(idx);
     // each segment is the index of a leaf.
     size_t seg_idx = thrust::get<0>(d_keys[i]);
-    auto seg = d_segments[seg_idx];
+    assert(seg_idx < d_indptr.size());
+    auto seg = d_indptr[seg_idx];
 
     double x = alpha * static_cast<double>(seg.n + 1);
     double k = std::floor(x) - 1;
     double d = (x - 1) - k;
 
     if (i == seg.begin) {
+      assert(d_row_index[d_sorted_idx[static_cast<size_t>(k)]] <= 8192);
       auto v0 = d_residue(d_row_index[d_sorted_idx[static_cast<size_t>(k)]], target_id);
       auto v1 = d_residue(d_row_index[d_sorted_idx[static_cast<size_t>(k) + 1]], target_id);
       d_results[seg_idx] = v0 + d * (v1 - v0);
@@ -832,7 +860,7 @@ class MeanAbsoluteError : public ObjFunction {
     return std::max(static_cast<size_t>(1), info.labels.Shape(1));
   }
 
-  struct ObjInfo Task() const override {
+  ObjInfo Task() const override {
     return {ObjInfo::kRegression, true, true};
   }
 
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 383c8b90c399..322b1f01a5e0 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -429,6 +429,8 @@ struct GPUHistMakerDevice {
     auto& segments = p_out_row_indices->back().indptr;
     auto const& tree = *p_tree;
     for (size_t nidx = 0; nidx < row_partitioner->ridx_segments_.size(); ++nidx) {
+      auto debug_seg = row_partitioner->ridx_segments_[nidx];
+      std::cout << "begin:" << debug_seg.begin << " size:"<< debug_seg.Size() << ", nidx:" << nidx << ", tree[nidx].IsLeaf(): " << tree[nidx].IsLeaf() << ", left:" << tree[nidx].LeftChild() << std::endl;
       if (tree[nidx].IsLeaf()) {
         auto const& seg = row_partitioner->ridx_segments_[nidx];
         // fixme: subsample
@@ -436,11 +438,20 @@ struct GPUHistMakerDevice {
             RowIndexCache::Segment{seg.begin, seg.Size(), static_cast<bst_node_t>(nidx)});
       }
     }
+    std::stable_sort(segments.begin(), segments.end(), [](auto l, auto r) {
+      return l.begin < r.begin;
+    });
     auto in = row_partitioner->GetRows();
     p_out_row_indices->back().row_index.Resize(in.size());
     auto d_row_index = p_out_row_indices->back().row_index.DeviceSpan();
-    dh::safe_cuda(
-        cudaMemcpyAsync(d_row_index.data(), in.data(), in.size_bytes(), cudaMemcpyDeviceToDevice));
+    thrust::copy(thrust::device, dh::tcbegin(in), dh::tcend(in), dh::tbegin(d_row_index));
+    // dh::DebugSyncDevice();
+    // std::cout << "GPU Hist" << std::endl;
+    // auto const& h_row_idx = p_out_row_indices->back().row_index.HostVector();
+    // for (auto idx : h_row_idx) {
+    //   std::cout << idx << ", ";
+    // }
+    // std::cout << std::endl;
   }
 
   void FinalisePositionInPage(EllpackPageImpl const *page,
@@ -740,6 +751,8 @@ class GPUHistMakerSpecialised {
     // rescale learning rate according to size of trees
     float lr = param_.learning_rate;
     param_.learning_rate = lr / trees.size();
+
+    row_set_collection_.clear();
     // build tree
     try {
       for (xgboost::RegTree* tree : trees) {
@@ -829,6 +842,10 @@ class GPUHistMakerSpecialised {
     return true;
   }
 
+  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const {
+    return row_set_collection_.at(tree_idx);
+  }
+
   TrainParam param_;   // NOLINT
   MetaInfo* info_{};   // NOLINT
 
@@ -916,7 +933,11 @@ class GPUHistMaker : public TreeUpdater {
   }
 
   common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
-    return row_set_collection_.at(tree_idx);
+    if (hist_maker_param_.single_precision_histogram) {
+      return float_maker_->GetRowIndexCache(tree_idx);
+    } else {
+      return double_maker_->GetRowIndexCache(tree_idx);
+    }
   }
 
   char const* Name() const override {
@@ -926,7 +947,6 @@ class GPUHistMaker : public TreeUpdater {
  private:
   GPUHistMakerTrainParam hist_maker_param_;
   ObjInfo task_;
-  std::vector<std::vector<RowIndexCache>> row_set_collection_;
   std::unique_ptr<GPUHistMakerSpecialised<GradientPair>> float_maker_;
   std::unique_ptr<GPUHistMakerSpecialised<GradientPairPrecise>> double_maker_;
 };

From 2a5f8bda23dfe90e246b509646d181f1ccb71b91 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Tue, 12 Apr 2022 18:57:44 +0800
Subject: [PATCH 023/124] rle.

---
 include/xgboost/tree_model.h          |  4 ++
 src/tree/gpu_hist/row_partitioner.cuh | 48 ++++++++++++++++++-
 src/tree/updater_gpu_hist.cu          | 68 ++++++++++++---------------
 3 files changed, 81 insertions(+), 39 deletions(-)

diff --git a/include/xgboost/tree_model.h b/include/xgboost/tree_model.h
index 0f39c39c3d33..bcac32adfe11 100644
--- a/include/xgboost/tree_model.h
+++ b/include/xgboost/tree_model.h
@@ -16,6 +16,7 @@
 #include <xgboost/feature_map.h>
 #include <xgboost/model.h>
 
+#include <cstddef>
 #include <limits>
 #include <vector>
 #include <string>
@@ -23,6 +24,7 @@
 #include <algorithm>
 #include <tuple>
 #include <stack>
+#include "xgboost/host_device_vector.h"
 
 namespace xgboost {
 
@@ -744,6 +746,8 @@ struct RowIndexCache {
 
   HostDeviceVector<size_t> row_index;
   std::vector<Segment> indptr;
+  HostDeviceVector<size_t> node_ptr;
+  HostDeviceVector<bst_node_t> node_idx;
 
   RowIndexCache(Context const* ctx, size_t n_samples) {
     if (!ctx->IsCPU()) {
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 5ff3c211f39c..12d863bf944a 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -2,8 +2,11 @@
  * Copyright 2017-2019 XGBoost contributors
  */
 #pragma once
+#include <cstddef>
 #include "xgboost/base.h"
 #include "../../common/device_helpers.cuh"
+#include "xgboost/generic_parameters.h"
+#include "xgboost/tree_model.h"
 
 namespace xgboost {
 namespace tree {
@@ -158,7 +161,8 @@ class RowPartitioner {
    * instance.
    */
   template <typename FinalisePositionOpT>
-  void FinalisePosition(FinalisePositionOpT op) {
+  void FinalisePosition(Context const* ctx, RegTree const* p_tree,
+                        std::vector<RowIndexCache>* p_out_row_indices, FinalisePositionOpT op) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
     dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
@@ -168,6 +172,48 @@ class RowPartitioner {
       if (new_position == kIgnoredTreePosition) return;
       d_position[idx] = new_position;
     });
+
+    dh::Iota(ridx_.CurrentSpan());
+    // copy position to buffer
+    size_t n_samples = position_.Size();
+    dh::XGBDeviceAllocator<char> alloc;
+    auto position = position_.Other();
+    dh::safe_cuda(cudaMemcpyAsync(position, d_position, position_.CurrentSpan().size_bytes(),
+                                  cudaMemcpyDeviceToDevice));
+    // sort row index according to node index
+    thrust::stable_sort_by_key(thrust::cuda::par(alloc), position, position + n_samples,
+                               ridx_.Current());
+
+    size_t n_leaf = p_tree->GetNumLeaves();
+    dh::device_vector<size_t> unique_out(n_leaf);
+    dh::device_vector<size_t> counts_out(n_leaf);
+    dh::TemporaryArray<size_t> num_runs_out(1);
+
+    size_t nbytes;
+    cub::DeviceRunLengthEncode::Encode(nullptr, nbytes, position, unique_out.data().get(),
+                                       counts_out.data().get(), num_runs_out.data().get(),
+                                       n_samples);
+    dh::TemporaryArray<char> temp(nbytes);
+    cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, position, unique_out.data().get(),
+                                       counts_out.data().get(), num_runs_out.data().get(),
+                                       n_samples);
+
+    dh::XGBCachingDeviceAllocator<char> caching;
+    thrust::inclusive_scan(thrust::cuda::par(caching), counts_out.begin(), counts_out.end(),
+                           counts_out.begin());
+
+    auto& row_indices = p_out_row_indices->back();
+    // copy node index (leaf index)
+    row_indices.node_idx.SetDevice(ctx->gpu_id);
+    row_indices.node_idx.Resize(n_leaf);
+    auto d_node_idx = row_indices.node_idx.DeviceSpan();
+    thrust::copy(thrust::device, unique_out.begin(), unique_out.end(), dh::tbegin(d_node_idx));
+    // copy node pointer
+    row_indices.node_ptr.SetDevice(ctx->gpu_id);
+    row_indices.node_ptr.Resize(n_leaf + 1, 0);
+    auto d_node_ptr = row_indices.node_ptr.DeviceSpan();
+    thrust::inclusive_scan(thrust::cuda::par(caching), counts_out.begin(), counts_out.end(),
+                           dh::tbegin(d_node_ptr) + 1);
   }
 
   /**
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 322b1f01a5e0..eac0afd65482 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -408,43 +408,32 @@ struct GPUHistMakerDevice {
       dh::CopyToD(categories_segments, &d_categories_segments);
     }
 
-    if (row_partitioner->GetRows().size() != p_fmat->Info().num_row_) {
-      row_partitioner.reset();  // Release the device memory first before reallocating
-      row_partitioner.reset(new RowPartitioner(ctx_->gpu_id, p_fmat->Info().num_row_));
-    }
-    if (page->n_rows == p_fmat->Info().num_row_) {
-      FinalisePositionInPage(page, dh::ToSpan(d_nodes),
-                             dh::ToSpan(d_split_types), dh::ToSpan(d_categories),
-                             dh::ToSpan(d_categories_segments));
-    } else {
-      for (auto& batch : p_fmat->GetBatches<EllpackPage>(batch_param)) {
-        FinalisePositionInPage(batch.Impl(), dh::ToSpan(d_nodes),
-                               dh::ToSpan(d_split_types), dh::ToSpan(d_categories),
-                               dh::ToSpan(d_categories_segments));
-      }
-    }
-
     CHECK(p_out_row_indices->empty());
     p_out_row_indices->push_back(RowIndexCache{ctx_, p_fmat->Info().num_row_});
-    auto& segments = p_out_row_indices->back().indptr;
-    auto const& tree = *p_tree;
-    for (size_t nidx = 0; nidx < row_partitioner->ridx_segments_.size(); ++nidx) {
-      auto debug_seg = row_partitioner->ridx_segments_[nidx];
-      std::cout << "begin:" << debug_seg.begin << " size:"<< debug_seg.Size() << ", nidx:" << nidx << ", tree[nidx].IsLeaf(): " << tree[nidx].IsLeaf() << ", left:" << tree[nidx].LeftChild() << std::endl;
-      if (tree[nidx].IsLeaf()) {
-        auto const& seg = row_partitioner->ridx_segments_[nidx];
-        // fixme: subsample
-        segments.push_back(
-            RowIndexCache::Segment{seg.begin, seg.Size(), static_cast<bst_node_t>(nidx)});
-      }
-    }
-    std::stable_sort(segments.begin(), segments.end(), [](auto l, auto r) {
-      return l.begin < r.begin;
-    });
-    auto in = row_partitioner->GetRows();
-    p_out_row_indices->back().row_index.Resize(in.size());
-    auto d_row_index = p_out_row_indices->back().row_index.DeviceSpan();
-    thrust::copy(thrust::device, dh::tcbegin(in), dh::tcend(in), dh::tbegin(d_row_index));
+    FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
+                           dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments),
+                           p_out_row_indices);
+
+
+    // auto& segments = p_out_row_indices->back().indptr;
+    // auto const& tree = *p_tree;
+    // for (size_t nidx = 0; nidx < row_partitioner->ridx_segments_.size(); ++nidx) {
+    //   auto debug_seg = row_partitioner->ridx_segments_[nidx];
+    //   std::cout << "begin:" << debug_seg.begin << " size:"<< debug_seg.Size() << ", nidx:" << nidx << ", tree[nidx].IsLeaf(): " << tree[nidx].IsLeaf() << ", left:" << tree[nidx].LeftChild() << std::endl;
+    //   if (tree[nidx].IsLeaf()) {
+    //     auto const& seg = row_partitioner->ridx_segments_[nidx];
+    //     // fixme: subsample
+    //     segments.push_back(
+    //         RowIndexCache::Segment{seg.begin, seg.Size(), static_cast<bst_node_t>(nidx)});
+    //   }
+    // }
+    // std::stable_sort(segments.begin(), segments.end(), [](auto l, auto r) {
+    //   return l.begin < r.begin;
+    // });
+    // auto in = row_partitioner->GetRows();
+    // p_out_row_indices->back().row_index.Resize(in.size());
+    // auto d_row_index = p_out_row_indices->back().row_index.DeviceSpan();
+    // thrust::copy(thrust::device, dh::tcbegin(in), dh::tcend(in), dh::tbegin(d_row_index));
     // dh::DebugSyncDevice();
     // std::cout << "GPU Hist" << std::endl;
     // auto const& h_row_idx = p_out_row_indices->back().row_index.HostVector();
@@ -455,13 +444,15 @@ struct GPUHistMakerDevice {
   }
 
   void FinalisePositionInPage(EllpackPageImpl const *page,
+                              RegTree const* p_tree,
                               const common::Span<RegTree::Node> d_nodes,
                               common::Span<FeatureType const> d_feature_types,
                               common::Span<uint32_t const> categories,
-                              common::Span<RegTree::Segment> categories_segments) {
+                              common::Span<RegTree::Segment> categories_segments,
+                              std::vector<RowIndexCache>* p_out_row_indices) {
     auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
     row_partitioner->FinalisePosition(
-        [=] __device__(size_t row_id, int position) {
+        ctx_, p_tree, p_out_row_indices, [=] __device__(size_t row_id, int position) {
           // What happens if user prune the tree?
           if (!d_matrix.IsInRange(row_id)) {
             return RowPartitioner::kIgnoredTreePosition;
@@ -692,7 +683,8 @@ struct GPUHistMakerDevice {
 
         int left_child_nidx = tree[candidate.nid].LeftChild();
         int right_child_nidx = tree[candidate.nid].RightChild();
-        // Only create child entries if needed
+        // Only create child entries if needed_
+        p_tree->GetNumLeaves();
         if (GPUExpandEntry::ChildIsValid(param, tree.GetDepth(left_child_nidx),
                                          num_leaves)) {
           monitor.Start("UpdatePosition");

From 9e7cdac4f0314fbdb5618d4b9367552cc7ae41f8 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Tue, 12 Apr 2022 19:16:44 +0800
Subject: [PATCH 024/124] use it in obj.

---
 src/objective/regression_obj.cu | 63 ++++++++-------------------------
 1 file changed, 14 insertions(+), 49 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index d66bbfa26525..0f992d4b4628 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -685,49 +685,18 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
     d_residue(sample_id, target_id) = y - d_predt[i];
   });
 
-  dh::device_vector<size_t> segment_idx(row_index.indptr.size() + 1, 0);
-  auto d_segment_idx = dh::ToSpan(segment_idx);
-  dh::device_vector<RowIndexCache::Segment> indptr(row_index.indptr);
-  for (auto const& seg : row_index.indptr) {
-    std::cout << seg.nidx << ", begin:" << seg.begin  << ", n:" << seg.n << std::endl;
-  }
-  std::cout << std::endl;
-
-  std::cout << "device segment" << std::endl;
-  for (size_t i = 0; i < indptr.size(); ++i) {
-    RowIndexCache::Segment seg = indptr[i];
-    std::cout << seg.nidx << ", begin:" << seg.begin  << ", n:" << seg.n << std::endl;
-  }
-  std::cout << std::endl;
-
-  auto d_indptr = dh::ToSpan(indptr);
-  dh::LaunchN(d_segment_idx.size(), [=] XGBOOST_DEVICE(size_t i) {
-    if (i == d_segment_idx.size() - 1) {
-      d_segment_idx[i] = d_indptr[i - 1].begin + d_indptr[i - 1].n;
-      printf("last: %lu, nidx: %d, %lu, %lu, i: %lu, \n", d_segment_idx[i], d_indptr[i - 1].nidx, d_indptr[i - 1].begin, d_indptr[i - i].n, i);
-      return;
-    }
-    d_segment_idx[i] = d_indptr[i].begin;
-  });
-
   dh::XGBDeviceAllocator<char> alloc;
   dh::device_vector<size_t> sorted_idx(d_labels.Shape(0));
   dh::Iota(dh::ToSpan(sorted_idx));
-  {
-    std::vector<size_t> h_segments(segment_idx.size());
-    thrust::copy(segment_idx.begin(), segment_idx.end(), h_segments.begin());
-    // for (auto idx : h_segments) {
-    //   std::cout << idx << ", ";
-    // }
-    // std::cout << std::endl;
-  }
+
   using Tup = thrust::tuple<size_t, float>;
+  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
   auto key_it = dh::MakeTransformIterator<Tup>(
       thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
         auto idx = linalg::UnravelIndex(i, d_labels.Shape());
         size_t sample_id = std::get<0>(idx);
         size_t target_id = std::get<1>(idx);
-        auto leaf_idx = dh::SegmentId(d_segment_idx, sample_id);
+        auto leaf_idx = dh::SegmentId(d_leaf_ptr, sample_id);
         auto residue = d_residue(sample_id, target_id);
         return thrust::make_tuple(leaf_idx, residue);
       });
@@ -744,16 +713,10 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
                              });
 
   quantiles->SetDevice(ctx->gpu_id);
-  quantiles->Resize(row_index.indptr.size());
+  quantiles->Resize(row_index.node_idx.Size());
   auto d_results = quantiles->DeviceSpan();
 
   auto d_row_index = row_index.row_index.ConstDeviceSpan();
-  std::cout << "Row index" << std::endl;
-  auto const& h_row_index = row_index.row_index.HostVector();
-  for (auto idx : h_row_index) {
-    std::cout << idx << ", ";
-  }
-  std::cout << std::endl;
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
   auto d_keys = dh::ToSpan(keys);
 
@@ -762,14 +725,14 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
     size_t target_id = std::get<1>(idx);
     // each segment is the index of a leaf.
     size_t seg_idx = thrust::get<0>(d_keys[i]);
-    assert(seg_idx < d_indptr.size());
-    auto seg = d_indptr[seg_idx];
+    size_t begin = d_leaf_ptr[seg_idx];
+    size_t n = d_leaf_ptr[seg_idx + 1] - d_leaf_ptr[seg_idx];
 
-    double x = alpha * static_cast<double>(seg.n + 1);
+    double x = alpha * static_cast<double>(n + 1);
     double k = std::floor(x) - 1;
     double d = (x - 1) - k;
 
-    if (i == seg.begin) {
+    if (i == begin) {
       assert(d_row_index[d_sorted_idx[static_cast<size_t>(k)]] <= 8192);
       auto v0 = d_residue(d_row_index[d_sorted_idx[static_cast<size_t>(k)]], target_id);
       auto v1 = d_residue(d_row_index[d_sorted_idx[static_cast<size_t>(k) + 1]], target_id);
@@ -791,11 +754,13 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
 
   auto const& h_results = results.HostVector();
   auto& tree = *p_tree;
-  for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
-    auto seg = row_index.front().indptr[i];
+  auto const& h_node_idx = row_index.front().node_idx.HostVector();
+  for (size_t i = 0; i < h_node_idx.size(); ++i) {
+    auto nidx = h_node_idx[i];
+    // auto seg = row_index.front().indptr[i];
     auto q = h_results[i];
-    CHECK(tree[seg.nidx].IsLeaf());
-    tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
+    CHECK(tree[nidx].IsLeaf());
+    tree[nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }
 

From b1be3fe6ecb40c630c4dd167b35e1fcdc5b90860 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Tue, 12 Apr 2022 23:26:36 +0800
Subject: [PATCH 025/124] fixes

---
 src/common/stats.h                    |  5 ++
 src/objective/regression_obj.cu       | 75 +++++++++++++++++++++------
 src/tree/gpu_hist/row_partitioner.cuh | 49 ++++++++++++-----
 src/tree/hist/evaluate_splits.h       |  4 --
 src/tree/updater_gpu_hist.cu          |  6 +++
 5 files changed, 106 insertions(+), 33 deletions(-)

diff --git a/src/common/stats.h b/src/common/stats.h
index dfb6feff2178..993d6b1f05bf 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -88,6 +88,11 @@ float Percentile(double alpha, Iter const& begin, Iter const& end) {
   std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
   std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
                    [&](size_t l, size_t r) { return *(begin + l) < *(begin + r); });
+  std::cout << "CPU" << std::endl;
+  for (auto v : sorted_idx) {
+    std::cout << v << ", ";
+  }
+  std::cout << std::endl;
 
   auto val = [&](size_t i) { return *(begin + sorted_idx[i]); };
   static_assert(std::is_same<decltype(val(0)), float>::value, "");
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 0f992d4b4628..5be476341009 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -673,6 +673,13 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
                          HostDeviceVector<float>* quantiles) {
   CHECK(alpha >= 0 && alpha <= 1);
 
+  // std::cout << "Row index" << std::endl;
+  // auto h_row_idx = row_index.row_index.HostVector();
+  // for (auto v : h_row_idx) {
+  //   std::cout << v << ", ";
+  // }
+  // std::cout << std::endl;
+
   auto d_predt = predt.ConstDeviceSpan();
   auto d_labels = info.labels.View(ctx->gpu_id);
   linalg::Tensor<float, 2> residue{d_labels.Shape(), ctx->gpu_id};
@@ -685,25 +692,34 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
     d_residue(sample_id, target_id) = y - d_predt[i];
   });
 
-  dh::XGBDeviceAllocator<char> alloc;
+  // auto const& h_predt = predt.HostVector();
+  // auto const& h_labels = info.labels.HostView();
+  // std::cout << std::endl;
+  // for (size_t i = 0; i < predt.Size(); ++i) {
+  //   std::cout << "l:" << h_labels(i) << ", p:" << h_predt[i] << std::endl;
+  // }
+  // std::cout << std::endl;
+
   dh::device_vector<size_t> sorted_idx(d_labels.Shape(0));
   dh::Iota(dh::ToSpan(sorted_idx));
 
   using Tup = thrust::tuple<size_t, float>;
   auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
+  auto d_row_index = row_index.row_index.ConstDeviceSpan();
   auto key_it = dh::MakeTransformIterator<Tup>(
       thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
         auto idx = linalg::UnravelIndex(i, d_labels.Shape());
         size_t sample_id = std::get<0>(idx);
         size_t target_id = std::get<1>(idx);
         auto leaf_idx = dh::SegmentId(d_leaf_ptr, sample_id);
-        auto residue = d_residue(sample_id, target_id);
+        auto residue = d_residue(d_row_index[sample_id], target_id);
         return thrust::make_tuple(leaf_idx, residue);
       });
   dh::device_vector<Tup> keys(residue.Size());
   dh::XGBCachingDeviceAllocator<char> caching;
   thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
 
+  dh::XGBDeviceAllocator<char> alloc;
   thrust::stable_sort_by_key(thrust::cuda::par(alloc), keys.begin(), keys.end(), sorted_idx.begin(),
                              [=] XGBOOST_DEVICE(Tup const& l, Tup const& r) {
                                if (thrust::get<0>(l) != thrust::get<0>(r)) {
@@ -712,32 +728,41 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
                                return thrust::get<1>(l) < thrust::get<1>(r);  // residue
                              });
 
+  // std::cout << "GPU" << std::endl;
+  // for (size_t i = 0; i < sorted_idx.size(); ++i) {
+  //   auto v = sorted_idx[i];
+  //   std::cout << v << ", ";
+  // }
+  // std::cout << std::endl;
+
   quantiles->SetDevice(ctx->gpu_id);
   quantiles->Resize(row_index.node_idx.Size());
   auto d_results = quantiles->DeviceSpan();
 
-  auto d_row_index = row_index.row_index.ConstDeviceSpan();
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
   auto d_keys = dh::ToSpan(keys);
 
-  dh::LaunchN(residue.Size(), [=] XGBOOST_DEVICE(size_t i) {
-    auto idx = linalg::UnravelIndex(i, d_labels.Shape());
-    size_t target_id = std::get<1>(idx);
+  dh::LaunchN(row_index.node_idx.Size(), [=] XGBOOST_DEVICE(size_t i) {
+    size_t target_id = 0;
     // each segment is the index of a leaf.
-    size_t seg_idx = thrust::get<0>(d_keys[i]);
+    size_t seg_idx = i;
     size_t begin = d_leaf_ptr[seg_idx];
-    size_t n = d_leaf_ptr[seg_idx + 1] - d_leaf_ptr[seg_idx];
+    size_t n = d_leaf_ptr[seg_idx + 1] - begin;
+
+    if (alpha <= (1 / (n + 1))) {
+      d_results[i] = d_residue(d_row_index[d_sorted_idx[0]]);
+    }
+    if (alpha >= (n / (n + 1))) {
+      d_results[i] = d_residue(d_row_index[d_sorted_idx[d_sorted_idx.size() - 1]]);
+    }
 
     double x = alpha * static_cast<double>(n + 1);
     double k = std::floor(x) - 1;
     double d = (x - 1) - k;
-
-    if (i == begin) {
-      assert(d_row_index[d_sorted_idx[static_cast<size_t>(k)]] <= 8192);
-      auto v0 = d_residue(d_row_index[d_sorted_idx[static_cast<size_t>(k)]], target_id);
-      auto v1 = d_residue(d_row_index[d_sorted_idx[static_cast<size_t>(k) + 1]], target_id);
-      d_results[seg_idx] = v0 + d * (v1 - v0);
-    }
+    auto v0 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k)]], target_id);
+    auto v1 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k) + 1]], target_id);
+    // printf("x: %f, k: %f, d: %f, v0: %f, v1: %f\n", x, k, d, v0, v1);
+    d_results[seg_idx] = v0 + d * (v1 - v0);
   });
 }
 
@@ -757,9 +782,9 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   auto const& h_node_idx = row_index.front().node_idx.HostVector();
   for (size_t i = 0; i < h_node_idx.size(); ++i) {
     auto nidx = h_node_idx[i];
-    // auto seg = row_index.front().indptr[i];
     auto q = h_results[i];
     CHECK(tree[nidx].IsLeaf());
+    // std::cout << "nidx:" << nidx << ", q:" << q << std::endl;
     tree[nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }
@@ -767,6 +792,20 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
                         MetaInfo const& info, HostDeviceVector<float> const& prediction,
                         uint32_t target, float alpha, RegTree* p_tree) {
+
+  // std::cout << std::endl;
+  // auto const& h_labels_dbg = info.labels.HostView();
+  // auto const& h_predt_dgb = prediction.HostVector();
+  // for (size_t i = 0; i < prediction.Size(); ++i) {
+  //   std::cout << "l:" << h_labels_dbg(i) << ", p:" << h_predt_dgb[i] << std::endl;
+  // }
+  // std::cout << "Row index" << std::endl;
+  // auto h_row_idx = row_index.front().row_index.HostVector();
+  // for (auto v : h_row_idx) {
+  //   std::cout << v << ", ";
+  // }
+  // std::cout << std::endl;
+
   auto& tree = *p_tree;
   std::vector<float> quantiles;
   for (auto const& part : row_index) {
@@ -775,14 +814,15 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
       auto const& seg = part.indptr[k];
       CHECK(tree[seg.nidx].IsLeaf());
       auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
-      float q{0};
       auto h_labels = info.labels.HostView().Slice(linalg::All(), target);
       auto const& h_prediction = prediction.ConstHostVector();
+
       auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
         auto row_idx = h_row_set[i];
         return h_labels(row_idx) - h_prediction[row_idx];
       });
 
+      float q{0};
       if (info.weights_.Empty()) {
         q = common::Percentile(alpha, iter, iter + h_row_set.size());
       } else {
@@ -813,6 +853,7 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     auto seg = row_index.front().indptr[i];
     auto q = quantiles[i];
     CHECK(tree[seg.nidx].IsLeaf());
+    // std::cout << "nidx:" << seg.nidx << ", q:" << q << std::endl;
     tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 12d863bf944a..559493b3f6d2 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -165,24 +165,53 @@ class RowPartitioner {
                         std::vector<RowIndexCache>* p_out_row_indices, FinalisePositionOpT op) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
+    auto d_ridx_b = ridx_.Other();
+    auto sorted_position = position_.Other();
     dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
       auto position = d_position[idx];
       RowIndexT ridx = d_ridx[idx];
       bst_node_t new_position = op(ridx, position);
       if (new_position == kIgnoredTreePosition) return;
       d_position[idx] = new_position;
+      // printf("d_pos: %lu, %d\n", idx, new_position);
+      sorted_position[ridx] = new_position;
     });
 
+    // {
+    //   std::vector<bst_node_t> h_position(position_.Size());
+    //   auto it = thrust::device_ptr<bst_node_t>(d_position);
+    //   thrust::copy(it, it + h_position.size(), h_position.begin());
+    //   for (size_t i = 0; i < h_position.size(); ++i) {
+    //     std::cout << "pos:" << h_position[i] << std::endl;
+    //   }
+    // }
+
     dh::Iota(ridx_.CurrentSpan());
     // copy position to buffer
     size_t n_samples = position_.Size();
     dh::XGBDeviceAllocator<char> alloc;
-    auto position = position_.Other();
-    dh::safe_cuda(cudaMemcpyAsync(position, d_position, position_.CurrentSpan().size_bytes(),
-                                  cudaMemcpyDeviceToDevice));
+    // dh::LaunchN(position_.Size(), [=]XGBOOST_DEVICE(size_t idx) {
+    //   auto ridx = d_ridx[idx];
+    //   sorted_position[ridx] = d_position[idx];
+    // });
+    // dh::safe_cuda(cudaMemcpyAsync(sorted_position, d_position, position_.CurrentSpan().size_bytes(),
+    //                               cudaMemcpyDeviceToDevice));
+    auto& row_indices = p_out_row_indices->back();
     // sort row index according to node index
-    thrust::stable_sort_by_key(thrust::cuda::par(alloc), position, position + n_samples,
-                               ridx_.Current());
+    row_indices.row_index.SetDevice(ctx->gpu_id);
+    row_indices.row_index.Resize(ridx_.Size());
+    dh::Iota(row_indices.row_index.DeviceSpan());
+    thrust::stable_sort_by_key(thrust::cuda::par(alloc), sorted_position, sorted_position + n_samples,
+                               row_indices.row_index.DevicePointer());
+
+    // auto const& h_row_idx = row_indices.row_index.HostVector();
+    // std::vector<bst_node_t> h_position(position_.Size());
+    // auto it = thrust::device_ptr<bst_node_t>(sorted_position);
+    // thrust::copy(it, it + h_position.size(), h_position.begin());
+    // for (size_t i = 0; i < h_position.size(); ++i) {
+    //   std::cout << h_row_idx[i] << " pos:" << h_position[i] << std::endl;
+    // }
+    // std::cout << std::endl;
 
     size_t n_leaf = p_tree->GetNumLeaves();
     dh::device_vector<size_t> unique_out(n_leaf);
@@ -190,25 +219,21 @@ class RowPartitioner {
     dh::TemporaryArray<size_t> num_runs_out(1);
 
     size_t nbytes;
-    cub::DeviceRunLengthEncode::Encode(nullptr, nbytes, position, unique_out.data().get(),
+    cub::DeviceRunLengthEncode::Encode(nullptr, nbytes, sorted_position, unique_out.data().get(),
                                        counts_out.data().get(), num_runs_out.data().get(),
                                        n_samples);
     dh::TemporaryArray<char> temp(nbytes);
-    cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, position, unique_out.data().get(),
+    cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, sorted_position, unique_out.data().get(),
                                        counts_out.data().get(), num_runs_out.data().get(),
                                        n_samples);
 
-    dh::XGBCachingDeviceAllocator<char> caching;
-    thrust::inclusive_scan(thrust::cuda::par(caching), counts_out.begin(), counts_out.end(),
-                           counts_out.begin());
-
-    auto& row_indices = p_out_row_indices->back();
     // copy node index (leaf index)
     row_indices.node_idx.SetDevice(ctx->gpu_id);
     row_indices.node_idx.Resize(n_leaf);
     auto d_node_idx = row_indices.node_idx.DeviceSpan();
     thrust::copy(thrust::device, unique_out.begin(), unique_out.end(), dh::tbegin(d_node_idx));
     // copy node pointer
+    dh::XGBCachingDeviceAllocator<char> caching;
     row_indices.node_ptr.SetDevice(ctx->gpu_id);
     row_indices.node_ptr.Resize(n_leaf + 1, 0);
     auto d_node_ptr = row_indices.node_ptr.DeviceSpan();
diff --git a/src/tree/hist/evaluate_splits.h b/src/tree/hist/evaluate_splits.h
index 5646bfc85497..5a2ef5abcc49 100644
--- a/src/tree/hist/evaluate_splits.h
+++ b/src/tree/hist/evaluate_splits.h
@@ -401,11 +401,7 @@ void UpdatePredictionCacheImpl(GenericParameter const *ctx, RegTree const *p_las
     common::ParallelFor2d(space, ctx->Threads(), [&](size_t nidx, common::Range1d r) {
       if (!tree[nidx].IsDeleted() && tree[nidx].IsLeaf()) {
         auto const &rowset = part[nidx];
-        auto const &stats = snode[nidx];
         auto leaf_value = tree[nidx].LeafValue();
-        // auto leaf_value =
-        //     evaluator.CalcWeight(nidx, param, GradStats{stats.stats}) * param.learning_rate;
-        // CHECK_EQ(leaf, leaf_value);
         for (const size_t *it = rowset.begin + r.begin(); it < rowset.begin + r.end(); ++it) {
           out_preds(*it) += leaf_value;
         }
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index eac0afd65482..c53c18059e1c 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -354,6 +354,7 @@ struct GPUHistMakerDevice {
   }
 
   void UpdatePosition(int nidx, RegTree* p_tree) {
+    // std::cout << "UpdatePosition:" << nidx << std::endl;
     RegTree::Node split_node = (*p_tree)[nidx];
     auto split_type = p_tree->NodeSplitType(nidx);
     auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
@@ -382,6 +383,7 @@ struct GPUHistMakerDevice {
               new_position = split_node.RightChild();
             }
           }
+          // printf("ridx: %d, pos: %d\n", ridx, new_position);
           return new_position;
         });
   }
@@ -455,6 +457,7 @@ struct GPUHistMakerDevice {
         ctx_, p_tree, p_out_row_indices, [=] __device__(size_t row_id, int position) {
           // What happens if user prune the tree?
           if (!d_matrix.IsInRange(row_id)) {
+            // printf("out\n");
             return RowPartitioner::kIgnoredTreePosition;
           }
           auto node = d_nodes[position];
@@ -472,6 +475,7 @@ struct GPUHistMakerDevice {
                                        categories_segments[position].size);
                 go_left = common::Decision<false>(node_cats, element, node.DefaultLeft());
               } else {
+                // printf("r: %lu, e: %f, s: %f\n", row_id, element, node.SplitCond());
                 go_left = element <= node.SplitCond();
               }
               if (go_left) {
@@ -482,6 +486,8 @@ struct GPUHistMakerDevice {
             }
             node = d_nodes[position];
           }
+
+          // printf("final ridx: %lu, pos: %d\n", row_id, position);
           return position;
         });
   }

From bf4c5b8c9097bb54d06c65710cc273e7f1343e0f Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Tue, 12 Apr 2022 23:43:28 +0800
Subject: [PATCH 026/124] fix.

---
 src/tree/gpu_hist/row_partitioner.cuh | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 559493b3f6d2..65d76994d6ab 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -165,7 +165,7 @@ class RowPartitioner {
                         std::vector<RowIndexCache>* p_out_row_indices, FinalisePositionOpT op) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
-    auto d_ridx_b = ridx_.Other();
+    // auto d_ridx_b = ridx_.Other();
     auto sorted_position = position_.Other();
     dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
       auto position = d_position[idx];
@@ -186,7 +186,6 @@ class RowPartitioner {
     //   }
     // }
 
-    dh::Iota(ridx_.CurrentSpan());
     // copy position to buffer
     size_t n_samples = position_.Size();
     dh::XGBDeviceAllocator<char> alloc;

From 34680ac0650dca196fef9a287f54cf019c3c65ae Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 13:42:19 +0800
Subject: [PATCH 027/124] Cleanup & fix.

---
 src/common/common.h             |  9 +++++++
 src/common/stats.h              |  7 +----
 src/metric/auc.cu               | 28 +++++++-------------
 src/objective/regression_obj.cu | 46 ++++-----------------------------
 4 files changed, 25 insertions(+), 65 deletions(-)

diff --git a/src/common/common.h b/src/common/common.h
index 877feba81553..f67cc676df16 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -208,6 +208,15 @@ struct OptionalWeights {
 
   XGBOOST_DEVICE float operator[](size_t i) const { return weights.empty() ? dft : weights[i]; }
 };
+
+
+/**
+ * Last index of a group in a CSR style of index pointer.
+ */
+template <typename Idx>
+XGBOOST_DEVICE size_t LastOf(size_t group, common::Span<Idx> indptr) {
+  return indptr[group + 1] - 1;
+}
 }  // namespace common
 }  // namespace xgboost
 #endif  // XGBOOST_COMMON_COMMON_H_
diff --git a/src/common/stats.h b/src/common/stats.h
index 993d6b1f05bf..2c61330ba768 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -88,11 +88,6 @@ float Percentile(double alpha, Iter const& begin, Iter const& end) {
   std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
   std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
                    [&](size_t l, size_t r) { return *(begin + l) < *(begin + r); });
-  std::cout << "CPU" << std::endl;
-  for (auto v : sorted_idx) {
-    std::cout << v << ", ";
-  }
-  std::cout << std::endl;
 
   auto val = [&](size_t i) { return *(begin + sorted_idx[i]); };
   static_assert(std::is_same<decltype(val(0)), float>::value, "");
@@ -103,7 +98,7 @@ float Percentile(double alpha, Iter const& begin, Iter const& end) {
   if (alpha >= (n / (n + 1))) {
     return val(sorted_idx.size() - 1);
   }
-
+  assert(n != 0 && "The number of rows in a leaf can not be zero.");
   double x = alpha * static_cast<double>((n + 1));
   double k = std::floor(x) - 1;
   CHECK_GE(k, 0);
diff --git a/src/metric/auc.cu b/src/metric/auc.cu
index be89c015c93d..52504ecbfbdf 100644
--- a/src/metric/auc.cu
+++ b/src/metric/auc.cu
@@ -201,14 +201,6 @@ void Transpose(common::Span<float const> in, common::Span<float> out, size_t m,
   });
 }
 
-/**
- * Last index of a group in a CSR style of index pointer.
- */
-template <typename Idx>
-XGBOOST_DEVICE size_t LastOf(size_t group, common::Span<Idx> indptr) {
-  return indptr[group + 1] - 1;
-}
-
 double ScaleClasses(common::Span<double> results,
                     common::Span<double> local_area, common::Span<double> fp,
                     common::Span<double> tp, common::Span<double> auc,
@@ -300,9 +292,9 @@ void SegmentedReduceAUC(common::Span<size_t const> d_unique_idx,
         double fp, tp, fp_prev, tp_prev;
         if (i == d_unique_class_ptr[class_id]) {
           // first item is ignored, we use this thread to calculate the last item
-          thrust::tie(fp, tp) = d_fptp[LastOf(class_id, d_class_ptr)];
+          thrust::tie(fp, tp) = d_fptp[common::LastOf(class_id, d_class_ptr)];
           thrust::tie(fp_prev, tp_prev) =
-              d_neg_pos[d_unique_idx[LastOf(class_id, d_unique_class_ptr)]];
+              d_neg_pos[d_unique_idx[common::LastOf(class_id, d_unique_class_ptr)]];
         } else {
           thrust::tie(fp, tp) = d_fptp[d_unique_idx[i] - 1];
           thrust::tie(fp_prev, tp_prev) = d_neg_pos[d_unique_idx[i - 1]];
@@ -413,10 +405,10 @@ double GPUMultiClassAUCOVR(common::Span<float const> predts,
     }
     uint32_t class_id = d_unique_idx[i] / n_samples;
     d_neg_pos[d_unique_idx[i]] = d_fptp[d_unique_idx[i] - 1];
-    if (i == LastOf(class_id, d_unique_class_ptr)) {
+    if (i == common::LastOf(class_id, d_unique_class_ptr)) {
       // last one needs to be included.
-      size_t last = d_unique_idx[LastOf(class_id, d_unique_class_ptr)];
-      d_neg_pos[LastOf(class_id, d_class_ptr)] = d_fptp[last - 1];
+      size_t last = d_unique_idx[common::LastOf(class_id, d_unique_class_ptr)];
+      d_neg_pos[common::LastOf(class_id, d_class_ptr)] = d_fptp[last - 1];
       return;
     }
   });
@@ -592,7 +584,7 @@ GPURankingAUC(common::Span<float const> predts, MetaInfo const &info,
         auto data_group_begin = d_group_ptr[group_id];
         size_t n_samples = d_group_ptr[group_id + 1] - data_group_begin;
         // last item of current group
-        if (item.idx == LastOf(group_id, d_threads_group_ptr)) {
+        if (item.idx == common::LastOf(group_id, d_threads_group_ptr)) {
           if (item.w > 0) {
             s_d_auc[group_id] = item.predt / item.w;
           } else {
@@ -797,10 +789,10 @@ GPURankingPRAUCImpl(common::Span<float const> predts, MetaInfo const &info,
     }
     auto group_idx = dh::SegmentId(d_group_ptr, d_unique_idx[i]);
     d_neg_pos[d_unique_idx[i]] = d_fptp[d_unique_idx[i] - 1];
-    if (i == LastOf(group_idx, d_unique_class_ptr)) {
+    if (i == common::LastOf(group_idx, d_unique_class_ptr)) {
       // last one needs to be included.
-      size_t last = d_unique_idx[LastOf(group_idx, d_unique_class_ptr)];
-      d_neg_pos[LastOf(group_idx, d_group_ptr)] = d_fptp[last - 1];
+      size_t last = d_unique_idx[common::LastOf(group_idx, d_unique_class_ptr)];
+      d_neg_pos[common::LastOf(group_idx, d_group_ptr)] = d_fptp[last - 1];
       return;
     }
   });
@@ -821,7 +813,7 @@ GPURankingPRAUCImpl(common::Span<float const> predts, MetaInfo const &info,
     auto it = dh::MakeTransformIterator<thrust::pair<double, uint32_t>>(
         thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t g) {
           double fp, tp;
-          thrust::tie(fp, tp) = d_fptp[LastOf(g, d_group_ptr)];
+          thrust::tie(fp, tp) = d_fptp[common::LastOf(g, d_group_ptr)];
           double area = fp * tp;
           auto n_documents = d_group_ptr[g + 1] - d_group_ptr[g];
           if (area > 0 && n_documents >= 2) {
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 5be476341009..665d1cae65f2 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -673,13 +673,6 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
                          HostDeviceVector<float>* quantiles) {
   CHECK(alpha >= 0 && alpha <= 1);
 
-  // std::cout << "Row index" << std::endl;
-  // auto h_row_idx = row_index.row_index.HostVector();
-  // for (auto v : h_row_idx) {
-  //   std::cout << v << ", ";
-  // }
-  // std::cout << std::endl;
-
   auto d_predt = predt.ConstDeviceSpan();
   auto d_labels = info.labels.View(ctx->gpu_id);
   linalg::Tensor<float, 2> residue{d_labels.Shape(), ctx->gpu_id};
@@ -692,14 +685,6 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
     d_residue(sample_id, target_id) = y - d_predt[i];
   });
 
-  // auto const& h_predt = predt.HostVector();
-  // auto const& h_labels = info.labels.HostView();
-  // std::cout << std::endl;
-  // for (size_t i = 0; i < predt.Size(); ++i) {
-  //   std::cout << "l:" << h_labels(i) << ", p:" << h_predt[i] << std::endl;
-  // }
-  // std::cout << std::endl;
-
   dh::device_vector<size_t> sorted_idx(d_labels.Shape(0));
   dh::Iota(dh::ToSpan(sorted_idx));
 
@@ -728,13 +713,6 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
                                return thrust::get<1>(l) < thrust::get<1>(r);  // residue
                              });
 
-  // std::cout << "GPU" << std::endl;
-  // for (size_t i = 0; i < sorted_idx.size(); ++i) {
-  //   auto v = sorted_idx[i];
-  //   std::cout << v << ", ";
-  // }
-  // std::cout << std::endl;
-
   quantiles->SetDevice(ctx->gpu_id);
   quantiles->Resize(row_index.node_idx.Size());
   auto d_results = quantiles->DeviceSpan();
@@ -747,13 +725,15 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
     // each segment is the index of a leaf.
     size_t seg_idx = i;
     size_t begin = d_leaf_ptr[seg_idx];
-    size_t n = d_leaf_ptr[seg_idx + 1] - begin;
+    auto n = static_cast<double>(d_leaf_ptr[seg_idx + 1] - begin);
 
     if (alpha <= (1 / (n + 1))) {
-      d_results[i] = d_residue(d_row_index[d_sorted_idx[0]]);
+      d_results[i] = d_residue(d_row_index[d_sorted_idx[begin]]);
+      return;
     }
     if (alpha >= (n / (n + 1))) {
-      d_results[i] = d_residue(d_row_index[d_sorted_idx[d_sorted_idx.size() - 1]]);
+      d_results[i] = d_residue(d_row_index[d_sorted_idx[common::LastOf(seg_idx, d_leaf_ptr)]]);
+      return;
     }
 
     double x = alpha * static_cast<double>(n + 1);
@@ -761,7 +741,6 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
     double d = (x - 1) - k;
     auto v0 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k)]], target_id);
     auto v1 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k) + 1]], target_id);
-    // printf("x: %f, k: %f, d: %f, v0: %f, v1: %f\n", x, k, d, v0, v1);
     d_results[seg_idx] = v0 + d * (v1 - v0);
   });
 }
@@ -784,7 +763,6 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
     auto nidx = h_node_idx[i];
     auto q = h_results[i];
     CHECK(tree[nidx].IsLeaf());
-    // std::cout << "nidx:" << nidx << ", q:" << q << std::endl;
     tree[nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }
@@ -792,20 +770,6 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
                         MetaInfo const& info, HostDeviceVector<float> const& prediction,
                         uint32_t target, float alpha, RegTree* p_tree) {
-
-  // std::cout << std::endl;
-  // auto const& h_labels_dbg = info.labels.HostView();
-  // auto const& h_predt_dgb = prediction.HostVector();
-  // for (size_t i = 0; i < prediction.Size(); ++i) {
-  //   std::cout << "l:" << h_labels_dbg(i) << ", p:" << h_predt_dgb[i] << std::endl;
-  // }
-  // std::cout << "Row index" << std::endl;
-  // auto h_row_idx = row_index.front().row_index.HostVector();
-  // for (auto v : h_row_idx) {
-  //   std::cout << v << ", ";
-  // }
-  // std::cout << std::endl;
-
   auto& tree = *p_tree;
   std::vector<float> quantiles;
   for (auto const& part : row_index) {

From 2df0a5001fb78857c1987017760eda4cff4c8283 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 14:11:54 +0800
Subject: [PATCH 028/124] Commented code.

---
 src/tree/gpu_hist/row_partitioner.cuh | 45 ++++++---------------------
 src/tree/hist/evaluate_splits.h       |  1 -
 src/tree/updater_gpu_hist.cu          | 28 -----------------
 3 files changed, 10 insertions(+), 64 deletions(-)

diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 65d76994d6ab..b17f3ee5a761 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -156,16 +156,15 @@ class RowPartitioner {
    * construction is complete. Does not update any other meta information in
    * this data structure, so should only be used at the end of training.
    *
-   * \param op          Device lambda. Should provide the row index  and current
-   * position as an argument and return the new position for this training
-   * instance.
+   * \param p_out_row_indices Row partitions for each leaf.
+   * \param op Device lambda. Should provide the row index and current position as an
+   * argument and return the new position for this training instance.
    */
   template <typename FinalisePositionOpT>
   void FinalisePosition(Context const* ctx, RegTree const* p_tree,
                         std::vector<RowIndexCache>* p_out_row_indices, FinalisePositionOpT op) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
-    // auto d_ridx_b = ridx_.Other();
     auto sorted_position = position_.Other();
     dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
       auto position = d_position[idx];
@@ -177,44 +176,20 @@ class RowPartitioner {
       sorted_position[ridx] = new_position;
     });
 
-    // {
-    //   std::vector<bst_node_t> h_position(position_.Size());
-    //   auto it = thrust::device_ptr<bst_node_t>(d_position);
-    //   thrust::copy(it, it + h_position.size(), h_position.begin());
-    //   for (size_t i = 0; i < h_position.size(); ++i) {
-    //     std::cout << "pos:" << h_position[i] << std::endl;
-    //   }
-    // }
-
     // copy position to buffer
     size_t n_samples = position_.Size();
     dh::XGBDeviceAllocator<char> alloc;
-    // dh::LaunchN(position_.Size(), [=]XGBOOST_DEVICE(size_t idx) {
-    //   auto ridx = d_ridx[idx];
-    //   sorted_position[ridx] = d_position[idx];
-    // });
-    // dh::safe_cuda(cudaMemcpyAsync(sorted_position, d_position, position_.CurrentSpan().size_bytes(),
-    //                               cudaMemcpyDeviceToDevice));
     auto& row_indices = p_out_row_indices->back();
     // sort row index according to node index
     row_indices.row_index.SetDevice(ctx->gpu_id);
     row_indices.row_index.Resize(ridx_.Size());
     dh::Iota(row_indices.row_index.DeviceSpan());
-    thrust::stable_sort_by_key(thrust::cuda::par(alloc), sorted_position, sorted_position + n_samples,
-                               row_indices.row_index.DevicePointer());
-
-    // auto const& h_row_idx = row_indices.row_index.HostVector();
-    // std::vector<bst_node_t> h_position(position_.Size());
-    // auto it = thrust::device_ptr<bst_node_t>(sorted_position);
-    // thrust::copy(it, it + h_position.size(), h_position.begin());
-    // for (size_t i = 0; i < h_position.size(); ++i) {
-    //   std::cout << h_row_idx[i] << " pos:" << h_position[i] << std::endl;
-    // }
-    // std::cout << std::endl;
+    thrust::stable_sort_by_key(thrust::cuda::par(alloc), sorted_position,
+                               sorted_position + n_samples, row_indices.row_index.DevicePointer());
 
     size_t n_leaf = p_tree->GetNumLeaves();
-    dh::device_vector<size_t> unique_out(n_leaf);
-    dh::device_vector<size_t> counts_out(n_leaf);
+    dh::caching_device_vector<size_t> unique_out(n_leaf);
+    dh::caching_device_vector<size_t> counts_out(n_leaf);
     dh::TemporaryArray<size_t> num_runs_out(1);
 
     size_t nbytes;
@@ -222,9 +197,9 @@ class RowPartitioner {
                                        counts_out.data().get(), num_runs_out.data().get(),
                                        n_samples);
     dh::TemporaryArray<char> temp(nbytes);
-    cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, sorted_position, unique_out.data().get(),
-                                       counts_out.data().get(), num_runs_out.data().get(),
-                                       n_samples);
+    cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, sorted_position,
+                                       unique_out.data().get(), counts_out.data().get(),
+                                       num_runs_out.data().get(), n_samples);
 
     // copy node index (leaf index)
     row_indices.node_idx.SetDevice(ctx->gpu_id);
diff --git a/src/tree/hist/evaluate_splits.h b/src/tree/hist/evaluate_splits.h
index 5a2ef5abcc49..4e445a0680e5 100644
--- a/src/tree/hist/evaluate_splits.h
+++ b/src/tree/hist/evaluate_splits.h
@@ -390,7 +390,6 @@ void UpdatePredictionCacheImpl(GenericParameter const *ctx, RegTree const *p_las
 
   CHECK(p_last_tree);
   auto const &tree = *p_last_tree;
-  auto const &snode = hist_evaluator.Stats();
   auto evaluator = hist_evaluator.Evaluator();
   CHECK_EQ(out_preds.DeviceIdx(), GenericParameter::kCpuId);
   size_t n_nodes = p_last_tree->GetNodes().size();
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index c53c18059e1c..cfe7598c661c 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -415,34 +415,6 @@ struct GPUHistMakerDevice {
     FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
                            dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments),
                            p_out_row_indices);
-
-
-    // auto& segments = p_out_row_indices->back().indptr;
-    // auto const& tree = *p_tree;
-    // for (size_t nidx = 0; nidx < row_partitioner->ridx_segments_.size(); ++nidx) {
-    //   auto debug_seg = row_partitioner->ridx_segments_[nidx];
-    //   std::cout << "begin:" << debug_seg.begin << " size:"<< debug_seg.Size() << ", nidx:" << nidx << ", tree[nidx].IsLeaf(): " << tree[nidx].IsLeaf() << ", left:" << tree[nidx].LeftChild() << std::endl;
-    //   if (tree[nidx].IsLeaf()) {
-    //     auto const& seg = row_partitioner->ridx_segments_[nidx];
-    //     // fixme: subsample
-    //     segments.push_back(
-    //         RowIndexCache::Segment{seg.begin, seg.Size(), static_cast<bst_node_t>(nidx)});
-    //   }
-    // }
-    // std::stable_sort(segments.begin(), segments.end(), [](auto l, auto r) {
-    //   return l.begin < r.begin;
-    // });
-    // auto in = row_partitioner->GetRows();
-    // p_out_row_indices->back().row_index.Resize(in.size());
-    // auto d_row_index = p_out_row_indices->back().row_index.DeviceSpan();
-    // thrust::copy(thrust::device, dh::tcbegin(in), dh::tcend(in), dh::tbegin(d_row_index));
-    // dh::DebugSyncDevice();
-    // std::cout << "GPU Hist" << std::endl;
-    // auto const& h_row_idx = p_out_row_indices->back().row_index.HostVector();
-    // for (auto idx : h_row_idx) {
-    //   std::cout << idx << ", ";
-    // }
-    // std::cout << std::endl;
   }
 
   void FinalisePositionInPage(EllpackPageImpl const *page,

From f0101a6c04c6efe3f6b398578173e129663dd754 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 14:23:21 +0800
Subject: [PATCH 029/124] Start working on weighted.

---
 src/common/stats.h | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/common/stats.h b/src/common/stats.h
index 2c61330ba768..9f6b626b71c9 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -109,29 +109,33 @@ float Percentile(double alpha, Iter const& begin, Iter const& end) {
   return v0 + d * (v1 - v0);
 }
 
-inline float WeightedPercentile(float quantile, common::Span<size_t const> row_set,
-                                linalg::VectorView<float const> labels,
-                                linalg::VectorView<float const> weights) {
-  std::vector<size_t> sorted_idx(row_set.size());
+template <typename Iter>
+float WeightedPercentile(double quantile, Iter begin, Iter end,
+                         linalg::VectorView<float const> weights) {
+  auto n = static_cast<double>(std::distance(begin, end));
+  std::vector<size_t> sorted_idx(n);
   std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
   std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
-                   [&](size_t i, size_t j) { return labels(row_set[i]) < labels(row_set[j]); });
-  std::vector<float> weighted_cdf(row_set.size());
+                   [&](size_t l, size_t r) { return *(begin + l) < *(begin + r); });
+
+  auto val = [&](size_t i) { return *(begin + sorted_idx[i]); };
+
+  std::vector<float> weighted_cdf(n);  // S_n
   weighted_cdf[0] = weights(row_set[sorted_idx[0]]);
-  for (size_t i = 1; i < row_set.size(); ++i) {
+  for (size_t i = 1; i < n; ++i) {
     weighted_cdf[i] = weighted_cdf[i - 1] + weights(row_set[sorted_idx[i]]);
   }
   float thresh = weighted_cdf.back() * quantile;
   size_t pos =
       std::upper_bound(weighted_cdf.cbegin(), weighted_cdf.cend(), thresh) - weighted_cdf.cbegin();
-  pos = std::min(pos, static_cast<size_t>(row_set.size() - 1));
-  if (pos == 0 || pos == static_cast<size_t>(row_set.size() - 1)) {
+  pos = std::min(pos, static_cast<size_t>(n - 1));
+  if (pos == 0 || pos == static_cast<size_t>(n - 1)) {
     return labels(row_set[sorted_idx[pos]]);
   }
   CHECK_GE(thresh, weighted_cdf[pos - 1]);
   CHECK_LT(thresh, weighted_cdf[pos]);
-  float v1 = labels(row_set[sorted_idx[pos - 1]]);
-  float v2 = labels(row_set[sorted_idx[pos]]);
+  float v1 = val(pos - 1);
+  float v2 = val(pos);
   if (weighted_cdf[pos + 1] - weighted_cdf[pos] >= 1.0f) {
     return (thresh - weighted_cdf[pos]) / (weighted_cdf[pos + 1] - weighted_cdf[pos]) * (v2 - v2) +
            v1;

From 39045789fb72ab02818a8c40767fd91429c08071 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 15:02:36 +0800
Subject: [PATCH 030/124] Start working on weighted.

---
 include/xgboost/linalg.h        |  4 +-
 src/common/stats.cuh            | 90 ++++++++++++++++++++++++++++++
 src/common/stats.h              | 27 ++++-----
 src/objective/regression_obj.cu | 97 ++++++---------------------------
 4 files changed, 120 insertions(+), 98 deletions(-)
 create mode 100644 src/common/stats.cuh

diff --git a/include/xgboost/linalg.h b/include/xgboost/linalg.h
index 15b3a04235bf..47082b42d8f1 100644
--- a/include/xgboost/linalg.h
+++ b/include/xgboost/linalg.h
@@ -672,10 +672,10 @@ class Tensor {
    */
   template <typename I, int32_t D>
   explicit Tensor(I const (&shape)[D], int32_t device)
-      : Tensor{common::Span<I, D>{shape}, device} {}
+      : Tensor{common::Span<I const, D>{shape}, device} {}
 
   template <typename I, size_t D>
-  explicit Tensor(common::Span<I, D> shape, int32_t device) {
+  explicit Tensor(common::Span<I const, D> shape, int32_t device) {
     // No device unroll as this is a host only function.
     std::copy(shape.data(), shape.data() + D, shape_);
     for (auto i = D; i < kDim; ++i) {
diff --git a/src/common/stats.cuh b/src/common/stats.cuh
new file mode 100644
index 000000000000..19dc1a940288
--- /dev/null
+++ b/src/common/stats.cuh
@@ -0,0 +1,90 @@
+#ifndef XGBOOST_COMMON_STATS_CUH_
+#define XGBOOST_COMMON_STATS_CUH_
+
+#include <thrust/sort.h>
+#include "device_helpers.cuh"
+#include "linalg_op.cuh"
+#include "xgboost/generic_parameters.h"
+#include "xgboost/tree_model.h"
+
+namespace xgboost {
+namespace common {
+void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const& row_index,
+                         MetaInfo const& info, HostDeviceVector<float> const& predt,
+                         HostDeviceVector<float>* quantiles) {
+  CHECK(alpha >= 0 && alpha <= 1);
+
+  auto d_predt = predt.ConstDeviceSpan();
+  auto d_labels = info.labels.View(ctx->gpu_id);
+  linalg::Tensor<float, 2> residue{d_labels.Shape(), ctx->gpu_id};
+  auto d_residue = residue.View(ctx->gpu_id);
+  CHECK_EQ(d_predt.size(), d_labels.Size());
+  linalg::ElementWiseKernel(ctx, d_labels, [=] XGBOOST_DEVICE(size_t i, float y) mutable {
+    auto idx = linalg::UnravelIndex(i, d_labels.Shape());
+    size_t sample_id = std::get<0>(idx);
+    size_t target_id = std::get<1>(idx);
+    d_residue(sample_id, target_id) = y - d_predt[i];
+  });
+
+  dh::device_vector<size_t> sorted_idx(d_labels.Shape(0));
+  dh::Iota(dh::ToSpan(sorted_idx));
+
+  using Tup = thrust::tuple<size_t, float>;
+  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
+  auto d_row_index = row_index.row_index.ConstDeviceSpan();
+  auto key_it = dh::MakeTransformIterator<Tup>(
+      thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
+        auto idx = linalg::UnravelIndex(i, d_labels.Shape());
+        size_t sample_id = std::get<0>(idx);
+        size_t target_id = std::get<1>(idx);
+        auto leaf_idx = dh::SegmentId(d_leaf_ptr, sample_id);
+        auto residue = d_residue(d_row_index[sample_id], target_id);
+        return thrust::make_tuple(leaf_idx, residue);
+      });
+  dh::device_vector<Tup> keys(residue.Size());
+  dh::XGBCachingDeviceAllocator<char> caching;
+  thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
+
+  dh::XGBDeviceAllocator<char> alloc;
+  thrust::stable_sort_by_key(thrust::cuda::par(alloc), keys.begin(), keys.end(), sorted_idx.begin(),
+                             [=] XGBOOST_DEVICE(Tup const& l, Tup const& r) {
+                               if (thrust::get<0>(l) != thrust::get<0>(r)) {
+                                 return thrust::get<0>(l) < thrust::get<0>(r);  // segment index
+                               }
+                               return thrust::get<1>(l) < thrust::get<1>(r);  // residue
+                             });
+
+  quantiles->SetDevice(ctx->gpu_id);
+  quantiles->Resize(row_index.node_idx.Size());
+  auto d_results = quantiles->DeviceSpan();
+
+  auto d_sorted_idx = dh::ToSpan(sorted_idx);
+  auto d_keys = dh::ToSpan(keys);
+
+  dh::LaunchN(row_index.node_idx.Size(), [=] XGBOOST_DEVICE(size_t i) {
+    size_t target_id = 0;
+    // each segment is the index of a leaf.
+    size_t seg_idx = i;
+    size_t begin = d_leaf_ptr[seg_idx];
+    auto n = static_cast<double>(d_leaf_ptr[seg_idx + 1] - begin);
+
+    if (alpha <= (1 / (n + 1))) {
+      d_results[i] = d_residue(d_row_index[d_sorted_idx[begin]]);
+      return;
+    }
+    if (alpha >= (n / (n + 1))) {
+      d_results[i] = d_residue(d_row_index[d_sorted_idx[common::LastOf(seg_idx, d_leaf_ptr)]]);
+      return;
+    }
+
+    double x = alpha * static_cast<double>(n + 1);
+    double k = std::floor(x) - 1;
+    double d = (x - 1) - k;
+    auto v0 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k)]], target_id);
+    auto v1 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k) + 1]], target_id);
+    d_results[seg_idx] = v0 + d * (v1 - v0);
+  });
+}
+}  // namespace common
+}  // namespace xgboost
+#endif  // XGBOOST_COMMON_STATS_CUH_
diff --git a/src/common/stats.h b/src/common/stats.h
index 9f6b626b71c9..ac7124a4a8a8 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -109,9 +109,8 @@ float Percentile(double alpha, Iter const& begin, Iter const& end) {
   return v0 + d * (v1 - v0);
 }
 
-template <typename Iter>
-float WeightedPercentile(double quantile, Iter begin, Iter end,
-                         linalg::VectorView<float const> weights) {
+template <typename Iter, typename WeightIter>
+float WeightedPercentile(double quantile, Iter begin, Iter end, WeightIter weights) {
   auto n = static_cast<double>(std::distance(begin, end));
   std::vector<size_t> sorted_idx(n);
   std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
@@ -121,23 +120,21 @@ float WeightedPercentile(double quantile, Iter begin, Iter end,
   auto val = [&](size_t i) { return *(begin + sorted_idx[i]); };
 
   std::vector<float> weighted_cdf(n);  // S_n
-  weighted_cdf[0] = weights(row_set[sorted_idx[0]]);
+  weighted_cdf[0] = *(weights + sorted_idx[0]);
   for (size_t i = 1; i < n; ++i) {
-    weighted_cdf[i] = weighted_cdf[i - 1] + weights(row_set[sorted_idx[i]]);
+    weighted_cdf[i] = weighted_cdf[i - 1] + *(weights + sorted_idx[i]);
   }
   float thresh = weighted_cdf.back() * quantile;
-  size_t pos =
+  size_t idx =
       std::upper_bound(weighted_cdf.cbegin(), weighted_cdf.cend(), thresh) - weighted_cdf.cbegin();
-  pos = std::min(pos, static_cast<size_t>(n - 1));
-  if (pos == 0 || pos == static_cast<size_t>(n - 1)) {
-    return labels(row_set[sorted_idx[pos]]);
+  idx = std::min(idx, static_cast<size_t>(n - 1));
+  if (idx == 0 || idx == static_cast<size_t>(n - 1)) {
+    return val(idx);
   }
-  CHECK_GE(thresh, weighted_cdf[pos - 1]);
-  CHECK_LT(thresh, weighted_cdf[pos]);
-  float v1 = val(pos - 1);
-  float v2 = val(pos);
-  if (weighted_cdf[pos + 1] - weighted_cdf[pos] >= 1.0f) {
-    return (thresh - weighted_cdf[pos]) / (weighted_cdf[pos + 1] - weighted_cdf[pos]) * (v2 - v2) +
+  float v1 = val(idx - 1);
+  float v2 = val(idx);
+  if (weighted_cdf[idx + 1] - weighted_cdf[idx] >= 1.0f) {
+    return (thresh - weighted_cdf[idx]) / (weighted_cdf[idx + 1] - weighted_cdf[idx]) * (v2 - v2) +
            v1;
   } else {
     return v2;
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 665d1cae65f2..eaa9ba6356f9 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -33,6 +33,8 @@
 
 #if defined(XGBOOST_USE_CUDA)
 #include "../common/linalg_op.cuh"
+#include "../common/device_helpers.cuh"
+#include "../common/stats.cuh"
 #endif  // defined(XGBOOST_USE_CUDA)
 
 namespace xgboost {
@@ -668,83 +670,7 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .describe("Tweedie regression for insurance data.")
 .set_body([]() { return new TweedieRegression(); });
 
-void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const& row_index,
-                         MetaInfo const& info, HostDeviceVector<float> const& predt,
-                         HostDeviceVector<float>* quantiles) {
-  CHECK(alpha >= 0 && alpha <= 1);
-
-  auto d_predt = predt.ConstDeviceSpan();
-  auto d_labels = info.labels.View(ctx->gpu_id);
-  linalg::Tensor<float, 2> residue{d_labels.Shape(), ctx->gpu_id};
-  auto d_residue = residue.View(ctx->gpu_id);
-  CHECK_EQ(d_predt.size(), d_labels.Size());
-  linalg::ElementWiseKernel(ctx, d_labels, [=] XGBOOST_DEVICE(size_t i, float y) mutable {
-    auto idx = linalg::UnravelIndex(i, d_labels.Shape());
-    size_t sample_id = std::get<0>(idx);
-    size_t target_id = std::get<1>(idx);
-    d_residue(sample_id, target_id) = y - d_predt[i];
-  });
-
-  dh::device_vector<size_t> sorted_idx(d_labels.Shape(0));
-  dh::Iota(dh::ToSpan(sorted_idx));
-
-  using Tup = thrust::tuple<size_t, float>;
-  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
-  auto d_row_index = row_index.row_index.ConstDeviceSpan();
-  auto key_it = dh::MakeTransformIterator<Tup>(
-      thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
-        auto idx = linalg::UnravelIndex(i, d_labels.Shape());
-        size_t sample_id = std::get<0>(idx);
-        size_t target_id = std::get<1>(idx);
-        auto leaf_idx = dh::SegmentId(d_leaf_ptr, sample_id);
-        auto residue = d_residue(d_row_index[sample_id], target_id);
-        return thrust::make_tuple(leaf_idx, residue);
-      });
-  dh::device_vector<Tup> keys(residue.Size());
-  dh::XGBCachingDeviceAllocator<char> caching;
-  thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
-
-  dh::XGBDeviceAllocator<char> alloc;
-  thrust::stable_sort_by_key(thrust::cuda::par(alloc), keys.begin(), keys.end(), sorted_idx.begin(),
-                             [=] XGBOOST_DEVICE(Tup const& l, Tup const& r) {
-                               if (thrust::get<0>(l) != thrust::get<0>(r)) {
-                                 return thrust::get<0>(l) < thrust::get<0>(r);  // segment index
-                               }
-                               return thrust::get<1>(l) < thrust::get<1>(r);  // residue
-                             });
-
-  quantiles->SetDevice(ctx->gpu_id);
-  quantiles->Resize(row_index.node_idx.Size());
-  auto d_results = quantiles->DeviceSpan();
-
-  auto d_sorted_idx = dh::ToSpan(sorted_idx);
-  auto d_keys = dh::ToSpan(keys);
-
-  dh::LaunchN(row_index.node_idx.Size(), [=] XGBOOST_DEVICE(size_t i) {
-    size_t target_id = 0;
-    // each segment is the index of a leaf.
-    size_t seg_idx = i;
-    size_t begin = d_leaf_ptr[seg_idx];
-    auto n = static_cast<double>(d_leaf_ptr[seg_idx + 1] - begin);
-
-    if (alpha <= (1 / (n + 1))) {
-      d_results[i] = d_residue(d_row_index[d_sorted_idx[begin]]);
-      return;
-    }
-    if (alpha >= (n / (n + 1))) {
-      d_results[i] = d_residue(d_row_index[d_sorted_idx[common::LastOf(seg_idx, d_leaf_ptr)]]);
-      return;
-    }
-
-    double x = alpha * static_cast<double>(n + 1);
-    double k = std::floor(x) - 1;
-    double d = (x - 1) - k;
-    auto v0 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k)]], target_id);
-    auto v1 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k) + 1]], target_id);
-    d_results[seg_idx] = v0 + d * (v1 - v0);
-  });
-}
-
+#if defined(XGBOOST_USE_CUDA)
 void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const> row_index,
                           MetaInfo const& info, HostDeviceVector<float> const& prediction,
                           uint32_t target, float alpha, RegTree* p_tree) {
@@ -754,7 +680,8 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   auto const& part = row_index.front();
 
   HostDeviceVector<float> results;
-  SegmentedPercentile(ctx, alpha, part, info, prediction, &results);
+
+  common::SegmentedPercentile(ctx, alpha, part, info, prediction, &results);
 
   auto const& h_results = results.HostVector();
   auto& tree = *p_tree;
@@ -766,6 +693,7 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
     tree[nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }
+#endif  // defined(XGBOOST_USE_CUDA)
 
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
                         MetaInfo const& info, HostDeviceVector<float> const& prediction,
@@ -780,19 +708,22 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
       auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
       auto h_labels = info.labels.HostView().Slice(linalg::All(), target);
       auto const& h_prediction = prediction.ConstHostVector();
+      auto h_weights = linalg::MakeVec(&info.weights_);
 
       auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
         auto row_idx = h_row_set[i];
         return h_labels(row_idx) - h_prediction[row_idx];
       });
+      auto w_it = common::MakeIndexTransformIter([&](size_t i) -> float {
+        auto row_idx = h_row_set[i];
+        return h_weights(row_idx);
+      });
 
       float q{0};
       if (info.weights_.Empty()) {
         q = common::Percentile(alpha, iter, iter + h_row_set.size());
       } else {
-        q = common::WeightedPercentile(alpha, h_row_set,
-                                       info.labels.HostView().Slice(linalg::All(), target),
-                                       linalg::MakeVec(&info.weights_));
+        q = common::WeightedPercentile(alpha, iter, iter + h_row_set.size(), w_it);
       }
       results.at(k) = q;
     });
@@ -865,7 +796,11 @@ class MeanAbsoluteError : public ObjFunction {
     if (ctx_->IsCPU()) {
       UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
     } else {
+#if defined(XGBOOST_USE_CUDA)
       UpdateTreeLeafDevice(ctx_, row_index, info, prediction, target, 0.5, p_tree);
+#else
+      common::AssertGPUSupport();
+#endif  //  defined(XGBOOST_USE_CUDA)
     }
   }
 

From 82af757f447b1b773050b111841344353a6a2252 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 15:11:12 +0800
Subject: [PATCH 031/124] Move.

---
 src/common/common.h | 61 +++++++++++++++++++++++++++++++++++++++++++++
 src/common/stats.h  | 52 --------------------------------------
 2 files changed, 61 insertions(+), 52 deletions(-)

diff --git a/src/common/common.h b/src/common/common.h
index f67cc676df16..6e3a8128eb3e 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -165,6 +165,67 @@ class Range {
   Iterator end_;
 };
 
+/**
+ * \brief Transform iterator that takes an index and calls transform operator.
+ *
+ *   This is CPU-only right now as taking host device function as operator complicates the
+ *   code.  For device side one can use `thrust::transform_iterator` instead.
+ */
+template <typename Fn>
+class IndexTransformIter {
+  size_t iter_{0};
+  Fn fn_;
+
+ public:
+  using iterator_category = std::random_access_iterator_tag;  // NOLINT
+  using value_type = std::result_of_t<Fn(size_t)>;            // NOLINT
+  using difference_type = detail::ptrdiff_t;                  // NOLINT
+  using reference = std::add_lvalue_reference_t<value_type>;  // NOLINT
+  using pointer = std::add_pointer_t<value_type>;             // NOLINT
+
+ public:
+  /**
+   * \param op Transform operator, takes a size_t index as input.
+   */
+  explicit IndexTransformIter(Fn &&op) : fn_{op} {}
+  IndexTransformIter(IndexTransformIter const &) = default;
+
+  value_type operator*() const { return fn_(iter_); }
+
+  auto operator-(IndexTransformIter const &that) const { return iter_ - that.iter_; }
+
+  IndexTransformIter &operator++() {
+    iter_++;
+    return *this;
+  }
+  IndexTransformIter operator++(int) {
+    auto ret = *this;
+    ++(*this);
+    return ret;
+  }
+  IndexTransformIter &operator+=(difference_type n) {
+    iter_ += n;
+    return *this;
+  }
+  IndexTransformIter &operator-=(difference_type n) {
+    (*this) += -n;
+    return *this;
+  }
+  IndexTransformIter operator+(difference_type n) const {
+    auto ret = *this;
+    return ret += n;
+  }
+  IndexTransformIter operator-(difference_type n) const {
+    auto ret = *this;
+    return ret -= n;
+  }
+};
+
+template <typename Fn>
+auto MakeIndexTransformIter(Fn&& fn) {
+  return IndexTransformIter<Fn>(std::forward<Fn>(fn));
+}
+
 int AllVisibleGPUs();
 
 inline void AssertGPUSupport() {
diff --git a/src/common/stats.h b/src/common/stats.h
index ac7124a4a8a8..ca07c3f48712 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -13,58 +13,6 @@
 namespace xgboost {
 namespace common {
 
-template <typename Fn>
-class IndexTransformIter {
-  size_t iter_{0};
-  Fn fn_;
-
- public:
-  using iterator_category = std::random_access_iterator_tag;  // NOLINT
-  using value_type = std::result_of_t<Fn(size_t)>;            // NOLINT
-  using difference_type = detail::ptrdiff_t;                  // NOLINT
-  using reference = std::add_lvalue_reference_t<value_type>;  // NOLINT
-  using pointer = std::add_pointer_t<value_type>;             // NOLINT
-
- public:
-  XGBOOST_DEVICE explicit IndexTransformIter(Fn&& fn) : fn_{fn} {}
-  IndexTransformIter(IndexTransformIter const&) = default;
-
-  value_type operator*() const { return fn_(iter_); }
-
-  XGBOOST_DEVICE auto operator-(IndexTransformIter const& that) const { return iter_ - that.iter_; }
-
-  XGBOOST_DEVICE IndexTransformIter& operator++() {
-    iter_++;
-    return *this;
-  }
-  XGBOOST_DEVICE IndexTransformIter operator++(int) {
-    auto ret = *this;
-    ++(*this);
-    return ret;
-  }
-  XGBOOST_DEVICE IndexTransformIter& operator+=(difference_type n) {
-    iter_ += n;
-    return *this;
-  }
-  XGBOOST_DEVICE IndexTransformIter& operator-=(difference_type n) {
-    (*this) += -n;
-    return *this;
-  }
-  XGBOOST_DEVICE IndexTransformIter operator+(difference_type n) const {
-    auto ret = *this;
-    return ret += n;
-  }
-  XGBOOST_DEVICE IndexTransformIter operator-(difference_type n) const {
-    auto ret = *this;
-    return ret -= n;
-  }
-};
-
-template <typename Fn>
-auto MakeIndexTransformIter(Fn&& fn) {
-  return IndexTransformIter<Fn>(std::forward<Fn>(fn));
-}
-
 /**
  * \brief Percentile with masked array using linear interpolation.
  *

From 508312552bdb0b047e6d50abe471c1396d5b2924 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 15:36:00 +0800
Subject: [PATCH 032/124] Refactor.

---
 src/common/stats.cuh            | 65 +++++++++++++++++++++++++++------
 src/objective/regression_obj.cu |  6 ++-
 2 files changed, 58 insertions(+), 13 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index 19dc1a940288..b9764fd2db56 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -1,22 +1,26 @@
+/*!
+ * Copyright 2022 by XGBoost Contributors
+ */
 #ifndef XGBOOST_COMMON_STATS_CUH_
 #define XGBOOST_COMMON_STATS_CUH_
 
 #include <thrust/sort.h>
+
 #include "device_helpers.cuh"
 #include "linalg_op.cuh"
 #include "xgboost/generic_parameters.h"
+#include "xgboost/linalg.h"
 #include "xgboost/tree_model.h"
 
 namespace xgboost {
 namespace common {
-void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const& row_index,
-                         MetaInfo const& info, HostDeviceVector<float> const& predt,
-                         HostDeviceVector<float>* quantiles) {
-  CHECK(alpha >= 0 && alpha <= 1);
+namespace detail {
+inline void ResidueDevice(Context const* ctx, linalg::TensorView<float const, 2> d_labels,
+                          common::Span<float const> d_predt, linalg::Tensor<float, 2>* p_residue) {
+  linalg::Tensor<float, 2>& residue = *p_residue;
+  residue.SetDevice(ctx->gpu_id);
+  residue.Reshape(d_labels.Shape());
 
-  auto d_predt = predt.ConstDeviceSpan();
-  auto d_labels = info.labels.View(ctx->gpu_id);
-  linalg::Tensor<float, 2> residue{d_labels.Shape(), ctx->gpu_id};
   auto d_residue = residue.View(ctx->gpu_id);
   CHECK_EQ(d_predt.size(), d_labels.Size());
   linalg::ElementWiseKernel(ctx, d_labels, [=] XGBOOST_DEVICE(size_t i, float y) mutable {
@@ -25,8 +29,13 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
     size_t target_id = std::get<1>(idx);
     d_residue(sample_id, target_id) = y - d_predt[i];
   });
+}
 
-  dh::device_vector<size_t> sorted_idx(d_labels.Shape(0));
+inline void SortLeafWeights(linalg::TensorView<float const, 2> d_residue,
+                            RowIndexCache const& row_index,
+                            dh::device_vector<size_t>* p_sorted_idx) {
+  auto& sorted_idx = *p_sorted_idx;
+  sorted_idx.resize(d_residue.Shape(0));
   dh::Iota(dh::ToSpan(sorted_idx));
 
   using Tup = thrust::tuple<size_t, float>;
@@ -34,14 +43,14 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
   auto d_row_index = row_index.row_index.ConstDeviceSpan();
   auto key_it = dh::MakeTransformIterator<Tup>(
       thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
-        auto idx = linalg::UnravelIndex(i, d_labels.Shape());
+        auto idx = linalg::UnravelIndex(i, d_residue.Shape());
         size_t sample_id = std::get<0>(idx);
         size_t target_id = std::get<1>(idx);
         auto leaf_idx = dh::SegmentId(d_leaf_ptr, sample_id);
         auto residue = d_residue(d_row_index[sample_id], target_id);
         return thrust::make_tuple(leaf_idx, residue);
       });
-  dh::device_vector<Tup> keys(residue.Size());
+  dh::device_vector<Tup> keys(d_residue.Size());
   dh::XGBCachingDeviceAllocator<char> caching;
   thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
 
@@ -53,13 +62,29 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
                                }
                                return thrust::get<1>(l) < thrust::get<1>(r);  // residue
                              });
+}
+}  // namespace detail
+
+inline void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const& row_index,
+                                MetaInfo const& info, HostDeviceVector<float> const& predt,
+                                HostDeviceVector<float>* quantiles) {
+  CHECK(alpha >= 0 && alpha <= 1);
+
+  auto d_predt = predt.ConstDeviceSpan();
+  auto d_labels = info.labels.View(ctx->gpu_id);
+  linalg::Tensor<float, 2> residue;
+  detail::ResidueDevice(ctx, d_labels, d_predt, &residue);
+  auto d_residue = residue.View(ctx->gpu_id);
+
+  dh::device_vector<size_t> sorted_idx;
+  detail::SortLeafWeights(d_residue, row_index, &sorted_idx);
 
   quantiles->SetDevice(ctx->gpu_id);
   quantiles->Resize(row_index.node_idx.Size());
   auto d_results = quantiles->DeviceSpan();
-
+  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
+  auto d_row_index = row_index.row_index.ConstDeviceSpan();
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
-  auto d_keys = dh::ToSpan(keys);
 
   dh::LaunchN(row_index.node_idx.Size(), [=] XGBOOST_DEVICE(size_t i) {
     size_t target_id = 0;
@@ -85,6 +110,22 @@ void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const&
     d_results[seg_idx] = v0 + d * (v1 - v0);
   });
 }
+
+inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
+                                      RowIndexCache const& row_index, MetaInfo const& info,
+                                      HostDeviceVector<float> const& predt,
+                                      HostDeviceVector<float>* quantiles) {
+  CHECK(alpha >= 0 && alpha <= 1);
+  auto d_predt = predt.ConstDeviceSpan();
+  auto d_labels = info.labels.View(ctx->gpu_id);
+  linalg::Tensor<float, 2> residue{d_labels.Shape(), ctx->gpu_id};
+  detail::ResidueDevice(ctx, d_labels, d_predt, &residue);
+  auto d_residue = residue.View(ctx->gpu_id);
+
+  dh::device_vector<size_t> sorted_idx;
+  detail::SortLeafWeights(d_residue, row_index, &sorted_idx);
+}
+
 }  // namespace common
 }  // namespace xgboost
 #endif  // XGBOOST_COMMON_STATS_CUH_
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index eaa9ba6356f9..60d13420c434 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -680,8 +680,12 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   auto const& part = row_index.front();
 
   HostDeviceVector<float> results;
+  if (info.weights_.Empty()) {
+    common::SegmentedPercentile(ctx, alpha, part, info, prediction, &results);
+  } else {
+    common::SegmentedWeightedQuantile(ctx, alpha, part, info, prediction, &results);
+  }
 
-  common::SegmentedPercentile(ctx, alpha, part, info, prediction, &results);
 
   auto const& h_results = results.HostVector();
   auto& tree = *p_tree;

From cc0defd54c63fa8f4645d34bd6dce0b32333abb4 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 16:21:16 +0800
Subject: [PATCH 033/124] GPU weighted.

---
 src/common/stats.cuh            | 86 ++++++++++++++++++++++++++++-----
 src/common/stats.h              | 24 ++++-----
 src/objective/regression_obj.cu |  2 -
 3 files changed, 86 insertions(+), 26 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index b9764fd2db56..900e01deabd1 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -5,6 +5,7 @@
 #define XGBOOST_COMMON_STATS_CUH_
 
 #include <thrust/sort.h>
+#include <cstddef>
 
 #include "device_helpers.cuh"
 #include "linalg_op.cuh"
@@ -15,11 +16,12 @@
 namespace xgboost {
 namespace common {
 namespace detail {
-inline void ResidueDevice(Context const* ctx, linalg::TensorView<float const, 2> d_labels,
-                          common::Span<float const> d_predt, linalg::Tensor<float, 2>* p_residue) {
+inline void ResidulesPredtY(Context const* ctx, linalg::TensorView<float const, 2> d_labels,
+                            common::Span<float const> d_predt,
+                            linalg::Tensor<float, 2>* p_residue) {
   linalg::Tensor<float, 2>& residue = *p_residue;
   residue.SetDevice(ctx->gpu_id);
-  residue.Reshape(d_labels.Shape());
+  residue.Reshape(d_labels.Shape(0), d_labels.Shape(1));
 
   auto d_residue = residue.View(ctx->gpu_id);
   CHECK_EQ(d_predt.size(), d_labels.Size());
@@ -31,14 +33,14 @@ inline void ResidueDevice(Context const* ctx, linalg::TensorView<float const, 2>
   });
 }
 
-inline void SortLeafWeights(linalg::TensorView<float const, 2> d_residue,
-                            RowIndexCache const& row_index,
-                            dh::device_vector<size_t>* p_sorted_idx) {
+template <typename Tup>
+inline void SortLeafRows(linalg::TensorView<float const, 2> d_residue,
+                         RowIndexCache const& row_index, dh::device_vector<size_t>* p_sorted_idx,
+                         dh::device_vector<Tup>* p_keys) {
   auto& sorted_idx = *p_sorted_idx;
   sorted_idx.resize(d_residue.Shape(0));
   dh::Iota(dh::ToSpan(sorted_idx));
 
-  using Tup = thrust::tuple<size_t, float>;
   auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
   auto d_row_index = row_index.row_index.ConstDeviceSpan();
   auto key_it = dh::MakeTransformIterator<Tup>(
@@ -50,7 +52,8 @@ inline void SortLeafWeights(linalg::TensorView<float const, 2> d_residue,
         auto residue = d_residue(d_row_index[sample_id], target_id);
         return thrust::make_tuple(leaf_idx, residue);
       });
-  dh::device_vector<Tup> keys(d_residue.Size());
+  dh::device_vector<Tup>& keys = *p_keys;
+  keys.resize(d_residue.Size());
   dh::XGBCachingDeviceAllocator<char> caching;
   thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
 
@@ -73,11 +76,13 @@ inline void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache
   auto d_predt = predt.ConstDeviceSpan();
   auto d_labels = info.labels.View(ctx->gpu_id);
   linalg::Tensor<float, 2> residue;
-  detail::ResidueDevice(ctx, d_labels, d_predt, &residue);
+  detail::ResidulesPredtY(ctx, d_labels, d_predt, &residue);
   auto d_residue = residue.View(ctx->gpu_id);
 
   dh::device_vector<size_t> sorted_idx;
-  detail::SortLeafWeights(d_residue, row_index, &sorted_idx);
+  using Tup = thrust::tuple<size_t, float>;
+  dh::device_vector<Tup> keys;
+  detail::SortLeafRows(d_residue, row_index, &sorted_idx, &keys);
 
   quantiles->SetDevice(ctx->gpu_id);
   quantiles->Resize(row_index.node_idx.Size());
@@ -119,13 +124,68 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
   auto d_predt = predt.ConstDeviceSpan();
   auto d_labels = info.labels.View(ctx->gpu_id);
   linalg::Tensor<float, 2> residue{d_labels.Shape(), ctx->gpu_id};
-  detail::ResidueDevice(ctx, d_labels, d_predt, &residue);
+  detail::ResidulesPredtY(ctx, d_labels, d_predt, &residue);
   auto d_residue = residue.View(ctx->gpu_id);
 
   dh::device_vector<size_t> sorted_idx;
-  detail::SortLeafWeights(d_residue, row_index, &sorted_idx);
-}
+  using Tup = thrust::tuple<size_t, float>;
+  dh::device_vector<Tup> keys;
+  detail::SortLeafRows(d_residue, row_index, &sorted_idx, &keys);
+  auto d_sorted_idx = dh::ToSpan(sorted_idx);
+  auto d_keys = dh::ToSpan(keys);
+
+  auto weights = info.weights_.ConstDeviceSpan();
+
+  dh::device_vector<float> weights_cdf(weights.size());
+  CHECK_EQ(weights.size(), d_labels.Shape(0));
 
+  dh::XGBCachingDeviceAllocator<char> caching;
+  Span<size_t const> node_ptr = row_index.node_ptr.ConstDeviceSpan();
+  // fixme: avoid this binary search by reusing key
+  auto scan_key = dh::MakeTransformIterator<size_t>(
+      thrust::make_counting_iterator(0ul),
+      [=] XGBOOST_DEVICE(size_t i) { return thrust::get<0>(d_keys[i]); });
+  auto scan_val = dh::MakeTransformIterator<float>(
+      thrust::make_counting_iterator(0ul),
+      [=] XGBOOST_DEVICE(size_t i) { return weights[d_sorted_idx[i]]; });
+  thrust::inclusive_scan_by_key(thrust::cuda::par(caching), scan_key, scan_key + weights.size(),
+                                dh::tcbegin(weights), weights_cdf.begin());
+
+  quantiles->SetDevice(ctx->gpu_id);
+  quantiles->Resize(row_index.node_idx.Size());
+  auto d_results = quantiles->DeviceSpan();
+  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
+  auto d_row_index = row_index.row_index.ConstDeviceSpan();
+  auto d_weight_cdf = dh::ToSpan(weights_cdf);
+
+  dh::LaunchN(row_index.node_idx.Size(), [=] XGBOOST_DEVICE(size_t i) {
+    size_t target_id = 0;
+    size_t seg_idx = i;
+    size_t begin = d_leaf_ptr[seg_idx];
+    auto n = static_cast<double>(d_leaf_ptr[seg_idx + 1] - begin);
+    auto leaf_cdf = d_weight_cdf.subspan(begin, n);
+    float thresh = leaf_cdf.back() * alpha;
+
+    size_t idx = thrust::lower_bound(thrust::seq, leaf_cdf.data(),
+                                     leaf_cdf.data() + leaf_cdf.size(), thresh) -
+                 leaf_cdf.data();
+    idx = std::min(idx, static_cast<size_t>(n - 1));
+    if (idx == 0 || idx == static_cast<size_t>(n - 1)) {
+      d_results[i] = d_residue(d_row_index[d_sorted_idx[idx + begin]], target_id);
+      return;
+    }
+
+    float v0 = d_residue(d_row_index[d_sorted_idx[idx + begin]], target_id);
+    float v1 = d_residue(d_row_index[d_sorted_idx[idx + begin + 1]], target_id);
+
+    if (leaf_cdf[idx + 1] - leaf_cdf[idx] >= 1.0f) {
+      auto v = (thresh - leaf_cdf[idx]) / (leaf_cdf[idx + 1] - leaf_cdf[idx]) * (v1 - v1) + v0;
+      d_results[i] = v;
+    } else {
+      d_results[i] = v1;
+    }
+  });
+}
 }  // namespace common
 }  // namespace xgboost
 #endif  // XGBOOST_COMMON_STATS_CUH_
diff --git a/src/common/stats.h b/src/common/stats.h
index ca07c3f48712..1300189e2067 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -3,6 +3,7 @@
  */
 #ifndef XGBOOST_COMMON_STATS_H_
 #define XGBOOST_COMMON_STATS_H_
+#include <algorithm>
 #include <iterator>
 #include <limits>
 #include <vector>
@@ -67,25 +68,26 @@ float WeightedPercentile(double quantile, Iter begin, Iter end, WeightIter weigh
 
   auto val = [&](size_t i) { return *(begin + sorted_idx[i]); };
 
-  std::vector<float> weighted_cdf(n);  // S_n
-  weighted_cdf[0] = *(weights + sorted_idx[0]);
+  std::vector<float> weight_cdf(n);  // S_n
+  // weighted cdf is sorted during construction
+  weight_cdf[0] = *(weights + sorted_idx[0]);
   for (size_t i = 1; i < n; ++i) {
-    weighted_cdf[i] = weighted_cdf[i - 1] + *(weights + sorted_idx[i]);
+    weight_cdf[i] = weight_cdf[i - 1] + *(weights + sorted_idx[i]);
   }
-  float thresh = weighted_cdf.back() * quantile;
+
+  float thresh = weight_cdf.back() * quantile;
   size_t idx =
-      std::upper_bound(weighted_cdf.cbegin(), weighted_cdf.cend(), thresh) - weighted_cdf.cbegin();
+      std::lower_bound(weight_cdf.cbegin(), weight_cdf.cend(), thresh) - weight_cdf.cbegin();
   idx = std::min(idx, static_cast<size_t>(n - 1));
   if (idx == 0 || idx == static_cast<size_t>(n - 1)) {
     return val(idx);
   }
-  float v1 = val(idx - 1);
-  float v2 = val(idx);
-  if (weighted_cdf[idx + 1] - weighted_cdf[idx] >= 1.0f) {
-    return (thresh - weighted_cdf[idx]) / (weighted_cdf[idx + 1] - weighted_cdf[idx]) * (v2 - v2) +
-           v1;
+  float v0 = val(idx);
+  float v1 = val(idx + 1);
+  if (weight_cdf[idx + 1] - weight_cdf[idx] >= 1.0f) {
+    return (thresh - weight_cdf[idx]) / (weight_cdf[idx + 1] - weight_cdf[idx]) * (v1 - v1) + v0;
   } else {
-    return v2;
+    return v1;
   }
 }
 }  // namespace common
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 60d13420c434..c71f2d86d18a 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -686,7 +686,6 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
     common::SegmentedWeightedQuantile(ctx, alpha, part, info, prediction, &results);
   }
 
-
   auto const& h_results = results.HostVector();
   auto& tree = *p_tree;
   auto const& h_node_idx = row_index.front().node_idx.HostVector();
@@ -752,7 +751,6 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     auto seg = row_index.front().indptr[i];
     auto q = quantiles[i];
     CHECK(tree[seg.nidx].IsLeaf());
-    // std::cout << "nidx:" << seg.nidx << ", q:" << q << std::endl;
     tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }

From c722b362fa1e5db0b22fb327ebb2e266a7db3db2 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 17:02:42 +0800
Subject: [PATCH 034/124] Fix.

---
 src/gbm/gbtree.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index d866705f9301..f23527c1c834 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -220,7 +220,7 @@ void CopyGradient(HostDeviceVector<GradientPair> const* in_gpair, int32_t n_thre
 void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const& predictions,
                             ObjFunction const* obj, size_t gidx,
                             std::vector<std::unique_ptr<RegTree>>* p_trees) {
-  if (!obj) {
+  if (!obj || !obj->Task().zero_hess) {
     return;
   }
   auto& trees = *p_trees;

From 7c5006168316324b052580ba7856b55e0bda9d35 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 17:10:13 +0800
Subject: [PATCH 035/124] Refactor.

---
 src/tree/updater_approx.cc | 56 ++++++++++++++++++++++----------------
 1 file changed, 32 insertions(+), 24 deletions(-)

diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index d51df8760fb8..f77aa97f7085 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -19,6 +19,7 @@
 #include "param.h"
 #include "xgboost/base.h"
 #include "xgboost/json.h"
+#include "xgboost/tree_model.h"
 #include "xgboost/tree_updater.h"
 
 namespace xgboost {
@@ -154,6 +155,36 @@ class GloablApproxBuilder {
     monitor_->Stop(__func__);
   }
 
+  void FinalisePosition(RegTree const &tree, common::Span<float> hess,
+                        std::vector<RowIndexCache> *p_out_row_indices) {
+    monitor_->Start(__func__);
+    CHECK(p_out_row_indices->empty());
+    for (auto const &part : partitioner_) {
+      auto const &row_set = part.Partitions();
+      p_out_row_indices->emplace_back(ctx_, row_set.Data()->size());
+      auto &h_row_index = p_out_row_indices->back().row_index.HostVector();
+
+      auto begin = row_set.Data()->data();
+      for (auto node : row_set) {
+        if (!node.begin) {
+          continue;
+        }
+        CHECK(node.begin && tree[node.node_id].IsLeaf()) << " Offending node idx:" << node.node_id;
+        size_t offset = node.begin - begin;
+        CHECK_LT(offset, row_set.Data()->size()) << node.node_id;
+        size_t k = offset;
+        for (auto idx = node.begin; idx != node.end; ++idx) {
+          if (hess[*idx] != 0.f) {
+            h_row_index[k++] = *idx;
+          }
+        }
+        auto seg = RowIndexCache::Segment{offset, k - offset, node.node_id};
+        p_out_row_indices->back().indptr.push_back(seg);
+      }
+    }
+    monitor_->Stop(__func__);
+  }
+
  public:
   explicit GloablApproxBuilder(TrainParam param, MetaInfo const &info, GenericParameter const *ctx,
                                std::shared_ptr<common::ColumnSampler> column_sampler, ObjInfo task,
@@ -232,30 +263,7 @@ class GloablApproxBuilder {
       expand_set = driver.Pop();
     }
 
-    CHECK(p_out_row_indices->empty());
-    for (auto const &part : partitioner_) {
-      auto const &row_set = part.Partitions();
-      p_out_row_indices->emplace_back(ctx_, p_fmat->Info().num_row_);
-      auto &h_row_index = p_out_row_indices->back().row_index.HostVector();
-
-      auto begin = row_set.Data()->data();
-      for (auto node : row_set) {
-        if (!node.begin) {
-          continue;
-        }
-        CHECK(node.begin && tree[node.node_id].IsLeaf()) << " Offending node idx:" << node.node_id;
-        size_t offset = node.begin - begin;
-        CHECK_LT(offset, p_fmat->Info().num_row_) << node.node_id;
-        size_t k = offset;
-        for (auto idx = node.begin; idx != node.end; ++idx) {
-          if (hess[*idx] != 0.f) {
-            h_row_index[k++] = *idx;
-          }
-        }
-        auto seg = RowIndexCache::Segment{offset, k - offset, node.node_id};
-        p_out_row_indices->back().indptr.push_back(seg);
-      }
-    }
+    this->FinalisePosition(tree, hess, p_out_row_indices);
   }
 };
 

From 055fa2f74a92a9b449b35e1c591406220517c9a5 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 17:35:38 +0800
Subject: [PATCH 036/124] hist.

---
 src/common/partition_builder.h    | 27 +++++++++++++++++++++++++
 src/tree/updater_approx.cc        | 28 ++++----------------------
 src/tree/updater_approx.h         |  8 ++++++++
 src/tree/updater_quantile_hist.cc | 33 +++++++++++++++++++++++--------
 src/tree/updater_quantile_hist.h  | 25 +++++++++++++++++++++--
 5 files changed, 87 insertions(+), 34 deletions(-)

diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index 3250b9d2bf25..8645929bdf6e 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -16,6 +16,7 @@
 
 #include "categorical.h"
 #include "column_matrix.h"
+#include "xgboost/generic_parameters.h"
 #include "xgboost/tree_model.h"
 
 namespace xgboost {
@@ -279,6 +280,32 @@ class PartitionBuilder {
     return blocks_offsets_[nid] + begin / BlockSize;
   }
 
+  // Copy row partitions into global cache for reuse in objective
+  template <typename Sampledp>
+  void LeafPartition(Context const* ctx, RegTree const& tree, RowSetCollection const& row_set,
+                     std::vector<RowIndexCache>* p_out_row_indices, Sampledp sampledp) const {
+    p_out_row_indices->emplace_back(ctx, row_set.Data()->size());
+    auto& h_row_index = p_out_row_indices->back().row_index.HostVector();
+
+    auto begin = row_set.Data()->data();
+    for (auto node : row_set) {
+      if (!node.begin) {
+        continue;
+      }
+      CHECK(node.begin && tree[node.node_id].IsLeaf());
+      size_t offset = node.begin - begin;
+      CHECK_LT(offset, row_set.Data()->size()) << node.node_id;
+      size_t k = offset;
+      for (auto idx = node.begin; idx != node.end; ++idx) {
+        if (!sampledp(*idx)) {
+          h_row_index[k++] = *idx;
+        }
+      }
+      auto seg = RowIndexCache::Segment{offset, k - offset, node.node_id};
+      p_out_row_indices->back().indptr.push_back(seg);
+    }
+  }
+
  protected:
   struct BlockInfo{
     size_t n_left;
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index f77aa97f7085..aa9de4f71858 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -155,32 +155,12 @@ class GloablApproxBuilder {
     monitor_->Stop(__func__);
   }
 
-  void FinalisePosition(RegTree const &tree, common::Span<float> hess,
-                        std::vector<RowIndexCache> *p_out_row_indices) {
+  void LeafPartition(RegTree const &tree, common::Span<float> hess,
+                     std::vector<RowIndexCache> *p_out_row_indices) {
     monitor_->Start(__func__);
     CHECK(p_out_row_indices->empty());
     for (auto const &part : partitioner_) {
-      auto const &row_set = part.Partitions();
-      p_out_row_indices->emplace_back(ctx_, row_set.Data()->size());
-      auto &h_row_index = p_out_row_indices->back().row_index.HostVector();
-
-      auto begin = row_set.Data()->data();
-      for (auto node : row_set) {
-        if (!node.begin) {
-          continue;
-        }
-        CHECK(node.begin && tree[node.node_id].IsLeaf()) << " Offending node idx:" << node.node_id;
-        size_t offset = node.begin - begin;
-        CHECK_LT(offset, row_set.Data()->size()) << node.node_id;
-        size_t k = offset;
-        for (auto idx = node.begin; idx != node.end; ++idx) {
-          if (hess[*idx] != 0.f) {
-            h_row_index[k++] = *idx;
-          }
-        }
-        auto seg = RowIndexCache::Segment{offset, k - offset, node.node_id};
-        p_out_row_indices->back().indptr.push_back(seg);
-      }
+      part.LeafPartition(ctx_, tree, hess, p_out_row_indices);
     }
     monitor_->Stop(__func__);
   }
@@ -263,7 +243,7 @@ class GloablApproxBuilder {
       expand_set = driver.Pop();
     }
 
-    this->FinalisePosition(tree, hess, p_out_row_indices);
+    this->LeafPartition(tree, hess, p_out_row_indices);
   }
 };
 
diff --git a/src/tree/updater_approx.h b/src/tree/updater_approx.h
index ec54da19e5b0..b662c6c4e81f 100644
--- a/src/tree/updater_approx.h
+++ b/src/tree/updater_approx.h
@@ -6,6 +6,7 @@
 #ifndef XGBOOST_TREE_UPDATER_APPROX_H_
 #define XGBOOST_TREE_UPDATER_APPROX_H_
 
+#include <cstddef>
 #include <limits>
 #include <utility>
 #include <vector>
@@ -18,6 +19,7 @@
 #include "hist/expand_entry.h"
 #include "hist/param.h"
 #include "param.h"
+#include "xgboost/generic_parameters.h"
 #include "xgboost/json.h"
 #include "xgboost/tree_updater.h"
 
@@ -122,6 +124,12 @@ class ApproxRowPartitioner {
 
   auto const &Partitions() const { return row_set_collection_; }
 
+  void LeafPartition(Context const *ctx, RegTree const &tree, common::Span<float const> hess,
+                     std::vector<RowIndexCache> *p_out_row_indices) const {
+    partition_builder_.LeafPartition(ctx, tree, this->Partitions(), p_out_row_indices,
+                                     [&](size_t idx) -> bool { return hess[idx] - .0f == .0f; });
+  }
+
   auto operator[](bst_node_t nidx) { return row_set_collection_[nidx]; }
   auto const &operator[](bst_node_t nidx) const { return row_set_collection_[nidx]; }
 
diff --git a/src/tree/updater_quantile_hist.cc b/src/tree/updater_quantile_hist.cc
index 0e1b6db47691..15aff5a16c01 100644
--- a/src/tree/updater_quantile_hist.cc
+++ b/src/tree/updater_quantile_hist.cc
@@ -53,11 +53,14 @@ void QuantileHistMaker::Update(HostDeviceVector<GradientPair> *gpair, DMatrix *d
     }
   }
 
+  row_set_collection_.clear();
   for (auto p_tree : trees) {
+    row_set_collection_.emplace_back();
+    auto &row_indices = row_set_collection_.back();
     if (hist_maker_param_.single_precision_histogram) {
-      this->float_builder_->UpdateTree(gpair, dmat, p_tree);
+      this->float_builder_->UpdateTree(gpair, dmat, p_tree, &row_indices);
     } else {
-      this->double_builder_->UpdateTree(gpair, dmat, p_tree);
+      this->double_builder_->UpdateTree(gpair, dmat, p_tree, &row_indices);
     }
   }
 
@@ -169,13 +172,27 @@ void QuantileHistMaker::Builder<GradientSumT>::BuildHistogram(
   }
 }
 
+template <typename GradientSumT>
+void QuantileHistMaker::Builder<GradientSumT>::LeafPartition(
+    RegTree const &tree, common::Span<GradientPair const> gpair,
+    std::vector<RowIndexCache> *p_out_row_indices) {
+  monitor_->Start(__func__);
+  CHECK(p_out_row_indices->empty());
+  for (auto const &part : partitioner_) {
+    part.LeafPartition(ctx_, tree, gpair, p_out_row_indices);
+  }
+  monitor_->Stop(__func__);
+}
+
 template <typename GradientSumT>
 void QuantileHistMaker::Builder<GradientSumT>::ExpandTree(
-    DMatrix *p_fmat, RegTree *p_tree, const std::vector<GradientPair> &gpair_h) {
+    DMatrix *p_fmat, RegTree *p_tree, const std::vector<GradientPair> &gpair_h,
+    std::vector<RowIndexCache> *p_out_row_indices) {
   monitor_->Start(__func__);
 
   Driver<CPUExpandEntry> driver(static_cast<TrainParam::TreeGrowPolicy>(param_.grow_policy));
   driver.Push(this->InitRoot(p_fmat, p_tree, gpair_h));
+  auto const &tree = *p_tree;
   bst_node_t num_leaves{1};
   auto expand_set = driver.Pop();
 
@@ -208,7 +225,6 @@ void QuantileHistMaker::Builder<GradientSumT>::ExpandTree(
     std::vector<CPUExpandEntry> best_splits;
     if (!valid_candidates.empty()) {
       this->BuildHistogram(p_fmat, p_tree, valid_candidates, gpair_h);
-      auto const &tree = *p_tree;
       for (auto const &candidate : valid_candidates) {
         int left_child_nidx = tree[candidate.nid].LeftChild();
         int right_child_nidx = tree[candidate.nid].RightChild();
@@ -228,12 +244,14 @@ void QuantileHistMaker::Builder<GradientSumT>::ExpandTree(
     expand_set = driver.Pop();
   }
 
+  this->LeafPartition(tree, gpair_h, p_out_row_indices);
   monitor_->Stop(__func__);
 }
 
 template <typename GradientSumT>
-void QuantileHistMaker::Builder<GradientSumT>::UpdateTree(HostDeviceVector<GradientPair> *gpair,
-                                                          DMatrix *p_fmat, RegTree *p_tree) {
+void QuantileHistMaker::Builder<GradientSumT>::UpdateTree(
+    HostDeviceVector<GradientPair> *gpair, DMatrix *p_fmat, RegTree *p_tree,
+    std::vector<RowIndexCache> *p_out_row_indices) {
   monitor_->Start(__func__);
 
   std::vector<GradientPair> *gpair_ptr = &(gpair->HostVector());
@@ -246,8 +264,7 @@ void QuantileHistMaker::Builder<GradientSumT>::UpdateTree(HostDeviceVector<Gradi
 
   this->InitData(p_fmat, *p_tree, gpair_ptr);
 
-  ExpandTree(p_fmat, p_tree, *gpair_ptr);
-
+  ExpandTree(p_fmat, p_tree, *gpair_ptr, p_out_row_indices);
   monitor_->Stop(__func__);
 }
 
diff --git a/src/tree/updater_quantile_hist.h b/src/tree/updater_quantile_hist.h
index 3c03a371ebfb..cd7ff1dab626 100644
--- a/src/tree/updater_quantile_hist.h
+++ b/src/tree/updater_quantile_hist.h
@@ -17,6 +17,7 @@
 #include <utility>
 #include <vector>
 
+#include "xgboost/base.h"
 #include "xgboost/data.h"
 #include "xgboost/json.h"
 
@@ -214,6 +215,15 @@ class HistRowPartitioner {
   size_t Size() const {
     return std::distance(row_set_collection_.begin(), row_set_collection_.end());
   }
+
+  void LeafPartition(Context const* ctx, RegTree const& tree,
+                     common::Span<GradientPair const> gpair,
+                     std::vector<RowIndexCache>* p_out_row_indices) const {
+    partition_builder_.LeafPartition(
+        ctx, tree, this->Partitions(), p_out_row_indices,
+        [&](size_t idx) -> bool { return gpair[idx].GetHess() - .0f == .0f; });
+  }
+
   auto& operator[](bst_node_t nidx) { return row_set_collection_[nidx]; }
   auto const& operator[](bst_node_t nidx) const { return row_set_collection_[nidx]; }
 };
@@ -266,6 +276,10 @@ class QuantileHistMaker: public TreeUpdater {
     return "grow_quantile_histmaker";
   }
 
+  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
+    return row_set_collection_.at(tree_idx);
+  }
+
  protected:
   CPUHistMakerTrainParam hist_maker_param_;
   // training parameter
@@ -289,7 +303,8 @@ class QuantileHistMaker: public TreeUpdater {
       monitor_->Init("Quantile::Builder");
     }
     // update one tree, growing
-    void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat, RegTree* p_tree);
+    void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat, RegTree* p_tree,
+                    std::vector<RowIndexCache>* p_out_row_indices);
 
     bool UpdatePredictionCache(DMatrix const* data, linalg::VectorView<float> out_preds) const;
 
@@ -308,7 +323,11 @@ class QuantileHistMaker: public TreeUpdater {
                         std::vector<CPUExpandEntry> const& valid_candidates,
                         std::vector<GradientPair> const& gpair);
 
-    void ExpandTree(DMatrix* p_fmat, RegTree* p_tree, const std::vector<GradientPair>& gpair_h);
+    void LeafPartition(RegTree const& tree, common::Span<GradientPair const> gpair,
+                       std::vector<RowIndexCache>* p_out_row_indices);
+
+    void ExpandTree(DMatrix* p_fmat, RegTree* p_tree, const std::vector<GradientPair>& gpair_h,
+                    std::vector<RowIndexCache>* p_out_row_indices);
 
    private:
     const size_t n_trees_;
@@ -334,6 +353,8 @@ class QuantileHistMaker: public TreeUpdater {
   };
 
  protected:
+  // cache for row partitions
+  std::vector<std::vector<RowIndexCache>> row_set_collection_;
   std::unique_ptr<Builder<float>> float_builder_;
   std::unique_ptr<Builder<double>> double_builder_;
   ObjInfo task_;

From 5ea25d018e69a9dedc57f36e77e3667457b62fe4 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 17:42:49 +0800
Subject: [PATCH 037/124] Format.

---
 src/common/stats.cuh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index 900e01deabd1..a55ac1cd9963 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -5,6 +5,7 @@
 #define XGBOOST_COMMON_STATS_CUH_
 
 #include <thrust/sort.h>
+
 #include <cstddef>
 
 #include "device_helpers.cuh"

From ce80cb31a8bd89f17521dcb305ca70dcca2ef40e Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 17:55:06 +0800
Subject: [PATCH 038/124] Cleanup.

---
 src/common/stats.cuh | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index a55ac1cd9963..49328a8290ae 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -17,6 +17,10 @@
 namespace xgboost {
 namespace common {
 namespace detail {
+/**
+ * \brief Compute the residue between label and prediction.  Can be simplifed once we have
+ *        prediction cache as matrix.
+ */
 inline void ResidulesPredtY(Context const* ctx, linalg::TensorView<float const, 2> d_labels,
                             common::Span<float const> d_predt,
                             linalg::Tensor<float, 2>* p_residue) {
@@ -34,6 +38,9 @@ inline void ResidulesPredtY(Context const* ctx, linalg::TensorView<float const,
   });
 }
 
+/**
+ * \brief Argsort based on residue value and row partition.
+ */
 template <typename Tup>
 inline void SortLeafRows(linalg::TensorView<float const, 2> d_residue,
                          RowIndexCache const& row_index, dh::device_vector<size_t>* p_sorted_idx,
@@ -133,6 +140,7 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
   dh::device_vector<Tup> keys;
   detail::SortLeafRows(d_residue, row_index, &sorted_idx, &keys);
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
+  auto d_row_index = row_index.row_index.ConstDeviceSpan();
   auto d_keys = dh::ToSpan(keys);
 
   auto weights = info.weights_.ConstDeviceSpan();
@@ -142,13 +150,12 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
 
   dh::XGBCachingDeviceAllocator<char> caching;
   Span<size_t const> node_ptr = row_index.node_ptr.ConstDeviceSpan();
-  // fixme: avoid this binary search by reusing key
   auto scan_key = dh::MakeTransformIterator<size_t>(
       thrust::make_counting_iterator(0ul),
       [=] XGBOOST_DEVICE(size_t i) { return thrust::get<0>(d_keys[i]); });
   auto scan_val = dh::MakeTransformIterator<float>(
       thrust::make_counting_iterator(0ul),
-      [=] XGBOOST_DEVICE(size_t i) { return weights[d_sorted_idx[i]]; });
+      [=] XGBOOST_DEVICE(size_t i) { return weights[d_row_index[d_sorted_idx[i]]]; });
   thrust::inclusive_scan_by_key(thrust::cuda::par(caching), scan_key, scan_key + weights.size(),
                                 dh::tcbegin(weights), weights_cdf.begin());
 
@@ -156,7 +163,6 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
   quantiles->Resize(row_index.node_idx.Size());
   auto d_results = quantiles->DeviceSpan();
   auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
-  auto d_row_index = row_index.row_index.ConstDeviceSpan();
   auto d_weight_cdf = dh::ToSpan(weights_cdf);
 
   dh::LaunchN(row_index.node_idx.Size(), [=] XGBOOST_DEVICE(size_t i) {
@@ -164,7 +170,8 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
     size_t seg_idx = i;
     size_t begin = d_leaf_ptr[seg_idx];
     auto n = static_cast<double>(d_leaf_ptr[seg_idx + 1] - begin);
-    auto leaf_cdf = d_weight_cdf.subspan(begin, n);
+    auto leaf_cdf = d_weight_cdf.subspan(begin, static_cast<size_t>(n));
+    auto leaf_sorted_idx = d_sorted_idx.subspan(begin, static_cast<size_t>(n));
     float thresh = leaf_cdf.back() * alpha;
 
     size_t idx = thrust::lower_bound(thrust::seq, leaf_cdf.data(),
@@ -172,12 +179,12 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
                  leaf_cdf.data();
     idx = std::min(idx, static_cast<size_t>(n - 1));
     if (idx == 0 || idx == static_cast<size_t>(n - 1)) {
-      d_results[i] = d_residue(d_row_index[d_sorted_idx[idx + begin]], target_id);
+      d_results[i] = d_residue(d_row_index[leaf_sorted_idx[idx]], target_id);
       return;
     }
 
-    float v0 = d_residue(d_row_index[d_sorted_idx[idx + begin]], target_id);
-    float v1 = d_residue(d_row_index[d_sorted_idx[idx + begin + 1]], target_id);
+    float v0 = d_residue(d_row_index[leaf_sorted_idx[idx]], target_id);
+    float v1 = d_residue(d_row_index[leaf_sorted_idx[idx + 1]], target_id);
 
     if (leaf_cdf[idx + 1] - leaf_cdf[idx] >= 1.0f) {
       auto v = (thresh - leaf_cdf[idx]) / (leaf_cdf[idx + 1] - leaf_cdf[idx]) * (v1 - v1) + v0;

From 6fbc58a168f1011648b0865761fbd5e74d4393b7 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 18:24:12 +0800
Subject: [PATCH 039/124] working on sampling.

---
 src/tree/gpu_hist/row_partitioner.cuh | 50 +++++++++++++++++++--------
 src/tree/updater_gpu_hist.cu          | 17 +++++----
 2 files changed, 46 insertions(+), 21 deletions(-)

diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index b17f3ee5a761..92f5b6cd04ca 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -3,9 +3,12 @@
  */
 #pragma once
 #include <cstddef>
+#include <cstdint>
+#include <limits>
 #include "xgboost/base.h"
 #include "../../common/device_helpers.cuh"
 #include "xgboost/generic_parameters.h"
+#include "xgboost/task.h"
 #include "xgboost/tree_model.h"
 
 namespace xgboost {
@@ -37,10 +40,10 @@ class RowPartitioner {
   using RowIndexT = bst_uint;
   struct Segment;
   static constexpr bst_node_t kIgnoredTreePosition = -1;
-  std::vector<Segment> ridx_segments_;
 
  private:
   int device_idx_;
+  std::vector<Segment> ridx_segments_;
   /*! \brief In here if you want to find the rows belong to a node nid, first you need to
    * get the indices segment from ridx_segments[nid], then get the row index that
    * represents position of row in input data X.  `RowPartitioner::GetRows` would be a
@@ -160,9 +163,10 @@ class RowPartitioner {
    * \param op Device lambda. Should provide the row index and current position as an
    * argument and return the new position for this training instance.
    */
-  template <typename FinalisePositionOpT>
-  void FinalisePosition(Context const* ctx, RegTree const* p_tree,
-                        std::vector<RowIndexCache>* p_out_row_indices, FinalisePositionOpT op) {
+  template <typename FinalisePositionOpT, typename Sampledp>
+  void FinalisePosition(Context const* ctx, RegTree const* p_tree, ObjInfo task,
+                        std::vector<RowIndexCache>* p_out_row_indices, FinalisePositionOpT op,
+                        Sampledp sampledp) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
     auto sorted_position = position_.Other();
@@ -170,12 +174,22 @@ class RowPartitioner {
       auto position = d_position[idx];
       RowIndexT ridx = d_ridx[idx];
       bst_node_t new_position = op(ridx, position);
-      if (new_position == kIgnoredTreePosition) return;
+      if (sampledp(ridx)) {
+        // push to the end
+        sorted_position[ridx] = std::numeric_limits<bst_node_t>::max();
+      } else {
+        sorted_position[ridx] = new_position;
+      }
+
+      if (new_position == kIgnoredTreePosition) {
+        return;
+      }
       d_position[idx] = new_position;
-      // printf("d_pos: %lu, %d\n", idx, new_position);
-      sorted_position[ridx] = new_position;
     });
 
+    if (!task.zero_hess) {
+      return;
+    }
     // copy position to buffer
     size_t n_samples = position_.Size();
     dh::XGBDeviceAllocator<char> alloc;
@@ -188,8 +202,10 @@ class RowPartitioner {
                                sorted_position + n_samples, row_indices.row_index.DevicePointer());
 
     size_t n_leaf = p_tree->GetNumLeaves();
-    dh::caching_device_vector<size_t> unique_out(n_leaf);
-    dh::caching_device_vector<size_t> counts_out(n_leaf);
+    // +1 for subsample, which is set to a unique value in above kernel.
+    size_t max_n_unique = n_leaf + 1;
+    dh::caching_device_vector<size_t> unique_out(max_n_unique);
+    dh::caching_device_vector<size_t> counts_out(max_n_unique);
     dh::TemporaryArray<size_t> num_runs_out(1);
 
     size_t nbytes;
@@ -201,18 +217,24 @@ class RowPartitioner {
                                        unique_out.data().get(), counts_out.data().get(),
                                        num_runs_out.data().get(), n_samples);
 
-    // copy node index (leaf index)
+    /**
+     * copy node index (leaf index)
+     */
     row_indices.node_idx.SetDevice(ctx->gpu_id);
     row_indices.node_idx.Resize(n_leaf);
     auto d_node_idx = row_indices.node_idx.DeviceSpan();
-    thrust::copy(thrust::device, unique_out.begin(), unique_out.end(), dh::tbegin(d_node_idx));
-    // copy node pointer
+    // don't copy the sampled values
+    thrust::copy(thrust::device, unique_out.begin(), unique_out.begin() + n_leaf,
+                 dh::tbegin(d_node_idx));
+    /**
+     * copy node pointer
+     */
     dh::XGBCachingDeviceAllocator<char> caching;
     row_indices.node_ptr.SetDevice(ctx->gpu_id);
     row_indices.node_ptr.Resize(n_leaf + 1, 0);
     auto d_node_ptr = row_indices.node_ptr.DeviceSpan();
-    thrust::inclusive_scan(thrust::cuda::par(caching), counts_out.begin(), counts_out.end(),
-                           dh::tbegin(d_node_ptr) + 1);
+    thrust::inclusive_scan(thrust::cuda::par(caching), counts_out.begin(),
+                           counts_out.begin() + n_leaf, dh::tbegin(d_node_ptr) + 1);
   }
 
   /**
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index cfe7598c661c..15cd3670305f 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -391,7 +391,8 @@ struct GPUHistMakerDevice {
   // After tree update is finished, update the position of all training
   // instances to their final leaf. This information is used later to update the
   // prediction cache
-  void FinalisePosition(RegTree const* p_tree, DMatrix* p_fmat, std::vector<RowIndexCache>* p_out_row_indices) {
+  void FinalisePosition(RegTree const* p_tree, DMatrix* p_fmat, ObjInfo task,
+                        std::vector<RowIndexCache>* p_out_row_indices) {
     dh::TemporaryArray<RegTree::Node> d_nodes(p_tree->GetNodes().size());
     dh::safe_cuda(cudaMemcpyAsync(d_nodes.data().get(), p_tree->GetNodes().data(),
                                   d_nodes.size() * sizeof(RegTree::Node),
@@ -413,7 +414,7 @@ struct GPUHistMakerDevice {
     CHECK(p_out_row_indices->empty());
     p_out_row_indices->push_back(RowIndexCache{ctx_, p_fmat->Info().num_row_});
     FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
-                           dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments),
+                           dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
                            p_out_row_indices);
   }
 
@@ -423,10 +424,13 @@ struct GPUHistMakerDevice {
                               common::Span<FeatureType const> d_feature_types,
                               common::Span<uint32_t const> categories,
                               common::Span<RegTree::Segment> categories_segments,
+                              ObjInfo task,
                               std::vector<RowIndexCache>* p_out_row_indices) {
     auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
+    auto d_gpair = this->gpair;
     row_partitioner->FinalisePosition(
-        ctx_, p_tree, p_out_row_indices, [=] __device__(size_t row_id, int position) {
+        ctx_, p_tree, task, p_out_row_indices,
+        [=] __device__(size_t row_id, int position) {
           // What happens if user prune the tree?
           if (!d_matrix.IsInRange(row_id)) {
             // printf("out\n");
@@ -447,7 +451,6 @@ struct GPUHistMakerDevice {
                                        categories_segments[position].size);
                 go_left = common::Decision<false>(node_cats, element, node.DefaultLeft());
               } else {
-                // printf("r: %lu, e: %f, s: %f\n", row_id, element, node.SplitCond());
                 go_left = element <= node.SplitCond();
               }
               if (go_left) {
@@ -459,9 +462,9 @@ struct GPUHistMakerDevice {
             node = d_nodes[position];
           }
 
-          // printf("final ridx: %lu, pos: %d\n", row_id, position);
           return position;
-        });
+        },
+        [d_gpair] __device__(size_t ridx) { return d_gpair[ridx].GetHess() - .0f == 0.f; });
   }
 
   void UpdatePredictionCache(linalg::VectorView<float> out_preds_d, RegTree const* p_tree) {
@@ -692,7 +695,7 @@ struct GPUHistMakerDevice {
     }
 
     monitor.Start("FinalisePosition");
-    this->FinalisePosition(p_tree, p_fmat, p_out_row_indices);
+    this->FinalisePosition(p_tree, p_fmat, task, p_out_row_indices);
     monitor.Stop("FinalisePosition");
   }
 };

From ebd500ce3abfc0a08bf8723bb3211d295d6ce2c3 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 19:06:54 +0800
Subject: [PATCH 040/124] Fix.

---
 src/common/stats.cuh            | 8 ++++----
 src/common/stats.h              | 3 ++-
 src/objective/regression_obj.cu | 1 +
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index 49328a8290ae..decc5a66a246 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -141,6 +141,7 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
   detail::SortLeafRows(d_residue, row_index, &sorted_idx, &keys);
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
   auto d_row_index = row_index.row_index.ConstDeviceSpan();
+  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
   auto d_keys = dh::ToSpan(keys);
 
   auto weights = info.weights_.ConstDeviceSpan();
@@ -152,17 +153,16 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
   Span<size_t const> node_ptr = row_index.node_ptr.ConstDeviceSpan();
   auto scan_key = dh::MakeTransformIterator<size_t>(
       thrust::make_counting_iterator(0ul),
-      [=] XGBOOST_DEVICE(size_t i) { return thrust::get<0>(d_keys[i]); });
+      [=] XGBOOST_DEVICE(size_t i) { return dh::SegmentId(d_leaf_ptr, i); });
   auto scan_val = dh::MakeTransformIterator<float>(
       thrust::make_counting_iterator(0ul),
       [=] XGBOOST_DEVICE(size_t i) { return weights[d_row_index[d_sorted_idx[i]]]; });
   thrust::inclusive_scan_by_key(thrust::cuda::par(caching), scan_key, scan_key + weights.size(),
-                                dh::tcbegin(weights), weights_cdf.begin());
+                                scan_val, weights_cdf.begin());
 
   quantiles->SetDevice(ctx->gpu_id);
   quantiles->Resize(row_index.node_idx.Size());
   auto d_results = quantiles->DeviceSpan();
-  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
   auto d_weight_cdf = dh::ToSpan(weights_cdf);
 
   dh::LaunchN(row_index.node_idx.Size(), [=] XGBOOST_DEVICE(size_t i) {
@@ -187,7 +187,7 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
     float v1 = d_residue(d_row_index[leaf_sorted_idx[idx + 1]], target_id);
 
     if (leaf_cdf[idx + 1] - leaf_cdf[idx] >= 1.0f) {
-      auto v = (thresh - leaf_cdf[idx]) / (leaf_cdf[idx + 1] - leaf_cdf[idx]) * (v1 - v1) + v0;
+      auto v = (thresh - leaf_cdf[idx]) / (leaf_cdf[idx + 1] - leaf_cdf[idx]) * (v1 - v0) + v0;
       d_results[i] = v;
     } else {
       d_results[i] = v1;
diff --git a/src/common/stats.h b/src/common/stats.h
index 1300189e2067..d6f7a3d75bd2 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -84,8 +84,9 @@ float WeightedPercentile(double quantile, Iter begin, Iter end, WeightIter weigh
   }
   float v0 = val(idx);
   float v1 = val(idx + 1);
+
   if (weight_cdf[idx + 1] - weight_cdf[idx] >= 1.0f) {
-    return (thresh - weight_cdf[idx]) / (weight_cdf[idx + 1] - weight_cdf[idx]) * (v1 - v1) + v0;
+    return (thresh - weight_cdf[idx]) / (weight_cdf[idx + 1] - weight_cdf[idx]) * (v1 - v0) + v0;
   } else {
     return v1;
   }
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index c71f2d86d18a..5c5056fbcb2a 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -751,6 +751,7 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     auto seg = row_index.front().indptr[i];
     auto q = quantiles[i];
     CHECK(tree[seg.nidx].IsLeaf());
+
     tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }

From ec82f8e5678a9fcad786999979d0d3c94938f3e1 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 22:58:30 +0800
Subject: [PATCH 041/124] Drop interpolation altogether.

---
 src/common/stats.h             | 20 ++++++---------
 tests/cpp/common/test_stats.cc | 47 +++++++++++++++++++++++++++++-----
 2 files changed, 49 insertions(+), 18 deletions(-)

diff --git a/src/common/stats.h b/src/common/stats.h
index d6f7a3d75bd2..de081a5b476e 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -58,6 +58,13 @@ float Percentile(double alpha, Iter const& begin, Iter const& end) {
   return v0 + d * (v1 - v0);
 }
 
+/**
+ * \brief Calculate the weighted quantile with step function. Unlike the unweighted
+ *        version, no interpolation is used.
+ *
+ *   See https://aakinshin.net/posts/weighted-quantiles/ for some discussion on computing
+ *   weighted quantile with interpolation.
+ */
 template <typename Iter, typename WeightIter>
 float WeightedPercentile(double quantile, Iter begin, Iter end, WeightIter weights) {
   auto n = static_cast<double>(std::distance(begin, end));
@@ -74,22 +81,11 @@ float WeightedPercentile(double quantile, Iter begin, Iter end, WeightIter weigh
   for (size_t i = 1; i < n; ++i) {
     weight_cdf[i] = weight_cdf[i - 1] + *(weights + sorted_idx[i]);
   }
-
   float thresh = weight_cdf.back() * quantile;
   size_t idx =
       std::lower_bound(weight_cdf.cbegin(), weight_cdf.cend(), thresh) - weight_cdf.cbegin();
   idx = std::min(idx, static_cast<size_t>(n - 1));
-  if (idx == 0 || idx == static_cast<size_t>(n - 1)) {
-    return val(idx);
-  }
-  float v0 = val(idx);
-  float v1 = val(idx + 1);
-
-  if (weight_cdf[idx + 1] - weight_cdf[idx] >= 1.0f) {
-    return (thresh - weight_cdf[idx]) / (weight_cdf[idx + 1] - weight_cdf[idx]) * (v1 - v0) + v0;
-  } else {
-    return v1;
-  }
+  return val(idx);
 }
 }  // namespace common
 }  // namespace xgboost
diff --git a/tests/cpp/common/test_stats.cc b/tests/cpp/common/test_stats.cc
index 327c709bde67..ee446ea6a049 100644
--- a/tests/cpp/common/test_stats.cc
+++ b/tests/cpp/common/test_stats.cc
@@ -1,3 +1,6 @@
+/*!
+ * Copyright 2022 by XGBoost Contributors
+ */
 #include <gtest/gtest.h>
 #include <xgboost/generic_parameters.h>
 
@@ -8,14 +11,46 @@ namespace common {
 TEST(Stats, Percentil) {
   linalg::Tensor<float, 1> arr({20.f, 0.f, 15.f, 50.f, 40.f, 0.f, 35.f}, {7}, Context::kCpuId);
   std::vector<size_t> index{0, 2, 3, 4, 6};
-  // auto percentile = Percentile(0.40f, Span<size_t const>{index}, arr.HostView());
-  // ASSERT_EQ(percentile, 26.0);
+  auto h_arr = arr.HostView();
+  auto beg = MakeIndexTransformIter([&](size_t i) { return h_arr(index[i]); });
+  auto end = beg + index.size();
+  auto percentile = Percentile(0.40f, beg, end);
+  ASSERT_EQ(percentile, 26.0);
 
-  // percentile = Percentile(0.20f, Span<size_t const>{index}, arr.HostView());
-  // ASSERT_EQ(percentile, 16.0);
+  percentile = Percentile(0.20f, beg, end);
+  ASSERT_EQ(percentile, 16.0);
 
-  // percentile = Percentile(0.10f, Span<size_t const>{index}, arr.HostView());
-  // ASSERT_EQ(percentile, 15.0);
+  percentile = Percentile(0.10f, beg, end);
+  ASSERT_EQ(percentile, 15.0);
+
+  {
+    std::vector<float> vec{1., 2., 3., 4., 5.};
+    auto beg = MakeIndexTransformIter([&](size_t i) { return vec[i]; });
+    auto end = beg + index.size();
+    auto percentile = Percentile(0.5f, beg, end);
+    ASSERT_EQ(percentile, 3.);
+  }
+}
+
+TEST(Stats, WeightedQuantile) {
+  linalg::Tensor<float, 1> arr({1.f, 2.f, 3.f, 4.f, 5.f}, {5}, Context::kCpuId);
+  linalg::Tensor<float, 1> weight({1.f, 1.f, 1.f, 1.f, 1.f}, {5}, Context::kCpuId);
+
+  auto h_arr = arr.HostView();
+  auto h_weight = weight.HostView();
+
+  auto beg = MakeIndexTransformIter([&](size_t i) { return h_arr(i); });
+  auto end = beg + arr.Size();
+  auto w = MakeIndexTransformIter([&](size_t i) { return h_weight(i); });
+
+  auto q = WeightedPercentile(0.50f, beg, end, w);
+  ASSERT_EQ(q, 3);
+
+  q = WeightedPercentile(0.0, beg, end, w);
+  ASSERT_EQ(q, 1);
+
+  q = WeightedPercentile(1.0, beg, end, w);
+  ASSERT_EQ(q, 5);
 }
 }  // namespace common
 }  // namespace xgboost

From 245e2026c08233cdc51913680712d0a0ebbfe5c6 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 13 Apr 2022 23:14:14 +0800
Subject: [PATCH 042/124] Drop on GPU.

---
 src/common/stats.cuh | 15 +--------------
 1 file changed, 1 insertion(+), 14 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index decc5a66a246..65f8e3b25850 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -178,20 +178,7 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
                                      leaf_cdf.data() + leaf_cdf.size(), thresh) -
                  leaf_cdf.data();
     idx = std::min(idx, static_cast<size_t>(n - 1));
-    if (idx == 0 || idx == static_cast<size_t>(n - 1)) {
-      d_results[i] = d_residue(d_row_index[leaf_sorted_idx[idx]], target_id);
-      return;
-    }
-
-    float v0 = d_residue(d_row_index[leaf_sorted_idx[idx]], target_id);
-    float v1 = d_residue(d_row_index[leaf_sorted_idx[idx + 1]], target_id);
-
-    if (leaf_cdf[idx + 1] - leaf_cdf[idx] >= 1.0f) {
-      auto v = (thresh - leaf_cdf[idx]) / (leaf_cdf[idx + 1] - leaf_cdf[idx]) * (v1 - v0) + v0;
-      d_results[i] = v;
-    } else {
-      d_results[i] = v1;
-    }
+    d_results[i] = d_residue(d_row_index[leaf_sorted_idx[idx]], target_id);
   });
 }
 }  // namespace common

From a0f97574bc698b01d3adcd4b4cd08a5edcbfe791 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 00:11:20 +0800
Subject: [PATCH 043/124] start poc for refactor.

---
 tests/cpp/common/test_stats.cu | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)
 create mode 100644 tests/cpp/common/test_stats.cu

diff --git a/tests/cpp/common/test_stats.cu b/tests/cpp/common/test_stats.cu
new file mode 100644
index 000000000000..777c8b9058db
--- /dev/null
+++ b/tests/cpp/common/test_stats.cu
@@ -0,0 +1,22 @@
+/*!
+ * Copyright 2022 by XGBoost Contributors
+ */
+#include <gtest/gtest.h>
+#include <cstddef>
+
+#include "../../../src/common/stats.cuh"
+#include "xgboost/host_device_vector.h"
+namespace xgboost {
+namespace common {
+TEST(Stats, GPUQuantile) {
+  linalg::Tensor<float, 1> arr({20.f, 0.f, 15.f, 50.f, 40.f, 0.f, 35.f}, {7}, Context::kCpuId);
+  HostDeviceVector<float> resutls;
+
+  auto d_arr = arr.View(0);
+  auto val_it = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
+                                                 [=](size_t i) { return d_arr(i); });
+}
+
+TEST(Stats, GPUWeightedQuantile) {}
+}  // namespace common
+}  // namespace xgboost

From 925e93c198cf34c4c2383f93b0428bd255461879 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 10:16:31 +0800
Subject: [PATCH 044/124] Use iterator for non-weight quantile.

---
 src/common/common.h             |  5 +++
 src/common/stats.cuh            | 79 +++++++++++++++++++++++----------
 src/objective/regression_obj.cu | 44 ++++++++++++++++--
 3 files changed, 101 insertions(+), 27 deletions(-)

diff --git a/src/common/common.h b/src/common/common.h
index 6e3a8128eb3e..95b7e4103261 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -278,6 +278,11 @@ template <typename Idx>
 XGBOOST_DEVICE size_t LastOf(size_t group, common::Span<Idx> indptr) {
   return indptr[group + 1] - 1;
 }
+
+template <typename Indexable>
+XGBOOST_DEVICE size_t LastOf(size_t group, Indexable const &indptr) {
+  return indptr[group + 1] - 1;
+}
 }  // namespace common
 }  // namespace xgboost
 #endif  // XGBOOST_COMMON_COMMON_H_
diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index 65f8e3b25850..fba291c7e6cd 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -4,9 +4,10 @@
 #ifndef XGBOOST_COMMON_STATS_CUH_
 #define XGBOOST_COMMON_STATS_CUH_
 
+#include <thrust/iterator/counting_iterator.h>
 #include <thrust/sort.h>
 
-#include <cstddef>
+#include <iterator>  // std::distance
 
 #include "device_helpers.cuh"
 #include "linalg_op.cuh"
@@ -74,52 +75,84 @@ inline void SortLeafRows(linalg::TensorView<float const, 2> d_residue,
                                return thrust::get<1>(l) < thrust::get<1>(r);  // residue
                              });
 }
+
+template <typename KeyIt, typename ValIt>
+inline void SortOnSegmented(KeyIt key_begin, KeyIt key_end, ValIt val_begin, ValIt val_end,
+                            dh::device_vector<size_t>* p_sorted_idx) {
+  using Tup = thrust::tuple<size_t, float>;
+  auto& sorted_idx = *p_sorted_idx;
+  size_t n = std::distance(val_begin, val_end);
+  sorted_idx.resize(n);
+  dh::Iota(dh::ToSpan(sorted_idx));
+  dh::device_vector<Tup> keys(sorted_idx.size());
+  auto key_it = dh::MakeTransformIterator<Tup>(
+      thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
+        auto leaf_idx = dh::SegmentId(key_begin, key_end, i);
+        auto residue = val_begin[i];
+        return thrust::make_tuple(leaf_idx, residue);
+      });
+  dh::XGBCachingDeviceAllocator<char> caching;
+  thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
+
+  dh::XGBDeviceAllocator<char> alloc;
+  thrust::stable_sort_by_key(thrust::cuda::par(alloc), keys.begin(), keys.end(), sorted_idx.begin(),
+                             [=] XGBOOST_DEVICE(Tup const& l, Tup const& r) {
+                               if (thrust::get<0>(l) != thrust::get<0>(r)) {
+                                 return thrust::get<0>(l) < thrust::get<0>(r);  // segment index
+                               }
+                               return thrust::get<1>(l) < thrust::get<1>(r);  // residue
+                             });
+}
 }  // namespace detail
 
-inline void SegmentedPercentile(Context const* ctx, double alpha, RowIndexCache const& row_index,
-                                MetaInfo const& info, HostDeviceVector<float> const& predt,
-                                HostDeviceVector<float>* quantiles) {
+/**
+ * \brief Compute segmented quantile on GPU.
+ *
+ * \tparam PtrIt Iterator for CSR style segments indptr
+ * \tparam ValIt Iterator for values
+ *
+ * \param alpha The p^th quantile we want to compute
+ *
+ *    std::distance(ptr_begin, ptr_end) should be equal to n_segments + 1
+ */
+template <typename PtrIt, typename ValIt>
+void SegmentedQuantile(Context const* ctx, double alpha, PtrIt key_begin, PtrIt key_end,
+                       ValIt val_begin, ValIt val_end, HostDeviceVector<float>* quantiles) {
   CHECK(alpha >= 0 && alpha <= 1);
 
-  auto d_predt = predt.ConstDeviceSpan();
-  auto d_labels = info.labels.View(ctx->gpu_id);
-  linalg::Tensor<float, 2> residue;
-  detail::ResidulesPredtY(ctx, d_labels, d_predt, &residue);
-  auto d_residue = residue.View(ctx->gpu_id);
-
   dh::device_vector<size_t> sorted_idx;
   using Tup = thrust::tuple<size_t, float>;
-  dh::device_vector<Tup> keys;
-  detail::SortLeafRows(d_residue, row_index, &sorted_idx, &keys);
+  detail::SortOnSegmented(key_begin, key_end, val_begin, val_end, &sorted_idx);
+  auto n_segments = std::distance(key_begin, key_end) - 1;
+  if (n_segments <= 0) {
+    return;
+  }
 
   quantiles->SetDevice(ctx->gpu_id);
-  quantiles->Resize(row_index.node_idx.Size());
+  quantiles->Resize(n_segments);
   auto d_results = quantiles->DeviceSpan();
-  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
-  auto d_row_index = row_index.row_index.ConstDeviceSpan();
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
 
-  dh::LaunchN(row_index.node_idx.Size(), [=] XGBOOST_DEVICE(size_t i) {
-    size_t target_id = 0;
+  dh::LaunchN(n_segments, [=] XGBOOST_DEVICE(size_t i) {
     // each segment is the index of a leaf.
     size_t seg_idx = i;
-    size_t begin = d_leaf_ptr[seg_idx];
-    auto n = static_cast<double>(d_leaf_ptr[seg_idx + 1] - begin);
+    size_t begin = key_begin[seg_idx];
+    auto n = static_cast<double>(key_begin[seg_idx + 1] - begin);
 
     if (alpha <= (1 / (n + 1))) {
-      d_results[i] = d_residue(d_row_index[d_sorted_idx[begin]]);
+      d_results[i] = val_begin[d_sorted_idx[begin]];
       return;
     }
     if (alpha >= (n / (n + 1))) {
-      d_results[i] = d_residue(d_row_index[d_sorted_idx[common::LastOf(seg_idx, d_leaf_ptr)]]);
+      d_results[i] = val_begin[common::LastOf(seg_idx, key_begin)];
       return;
     }
 
     double x = alpha * static_cast<double>(n + 1);
     double k = std::floor(x) - 1;
     double d = (x - 1) - k;
-    auto v0 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k)]], target_id);
-    auto v1 = d_residue(d_row_index[d_sorted_idx[begin + static_cast<size_t>(k) + 1]], target_id);
+    auto v0 = val_begin[d_sorted_idx[begin + static_cast<size_t>(k)]];
+    auto v1 = val_begin[d_sorted_idx[begin + static_cast<size_t>(k) + 1]];
     d_results[seg_idx] = v0 + d * (v1 - v0);
   });
 }
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 5c5056fbcb2a..34577b93dfdd 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -671,8 +671,30 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .set_body([]() { return new TweedieRegression(); });
 
 #if defined(XGBOOST_USE_CUDA)
+namespace detail {
+/**
+ * \brief Compute the residue between label and prediction.  Can be simplifed once we have
+ *        prediction cache as matrix.
+ */
+inline void ResidulesPredtY(Context const* ctx, linalg::TensorView<float const, 2> d_labels,
+                            common::Span<float const> d_predt,
+                            linalg::Tensor<float, 2>* p_residue) {
+  linalg::Tensor<float, 2>& residue = *p_residue;
+  residue.SetDevice(ctx->gpu_id);
+  residue.Reshape(d_labels.Shape(0), d_labels.Shape(1));
+
+  auto d_residue = residue.View(ctx->gpu_id);
+  CHECK_EQ(d_predt.size(), d_labels.Size());
+  linalg::ElementWiseKernel(ctx, d_labels, [=] XGBOOST_DEVICE(size_t i, float y) mutable {
+    auto idx = linalg::UnravelIndex(i, d_labels.Shape());
+    size_t sample_id = std::get<0>(idx);
+    size_t target_id = std::get<1>(idx);
+    d_residue(sample_id, target_id) = y - d_predt[i];
+  });
+}
+
 void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const> row_index,
-                          MetaInfo const& info, HostDeviceVector<float> const& prediction,
+                          MetaInfo const& info, HostDeviceVector<float> const& predt,
                           uint32_t target, float alpha, RegTree* p_tree) {
   dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
   CHECK_EQ(row_index.size(), 1)
@@ -680,10 +702,23 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   auto const& part = row_index.front();
 
   HostDeviceVector<float> results;
+  auto d_predt = predt.ConstDeviceSpan();
+  auto d_labels = info.labels.View(ctx->gpu_id);
+  linalg::Tensor<float, 2> residue;
+  ResidulesPredtY(ctx, d_labels, d_predt, &residue);
+  auto d_residue = residue.View(ctx->gpu_id);
+
+  auto d_row_index = part.row_index.ConstDeviceSpan();
+  auto key_it = part.node_ptr.ConstDeviceSpan().data();
+  auto val_it = dh::MakeTransformIterator<float>(
+      thrust::make_counting_iterator(0ul),
+      [=] XGBOOST_DEVICE(size_t i) { return d_residue(d_row_index[i]); });
+  CHECK_EQ(part.node_idx.Size() + 1, part.node_ptr.Size());
   if (info.weights_.Empty()) {
-    common::SegmentedPercentile(ctx, alpha, part, info, prediction, &results);
+    common::SegmentedQuantile(ctx, alpha, key_it, key_it + part.node_ptr.Size(), val_it,
+                              val_it + d_residue.Size(), &results);
   } else {
-    common::SegmentedWeightedQuantile(ctx, alpha, part, info, prediction, &results);
+    common::SegmentedWeightedQuantile(ctx, alpha, part, info, predt, &results);
   }
 
   auto const& h_results = results.HostVector();
@@ -696,6 +731,7 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
     tree[nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }
+}  // namespace detail
 #endif  // defined(XGBOOST_USE_CUDA)
 
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
@@ -800,7 +836,7 @@ class MeanAbsoluteError : public ObjFunction {
       UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
     } else {
 #if defined(XGBOOST_USE_CUDA)
-      UpdateTreeLeafDevice(ctx_, row_index, info, prediction, target, 0.5, p_tree);
+      detail::UpdateTreeLeafDevice(ctx_, row_index, info, prediction, target, 0.5, p_tree);
 #else
       common::AssertGPUSupport();
 #endif  //  defined(XGBOOST_USE_CUDA)

From 29bd9a6377594b55f2c09526de501b77183bf900 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 10:24:17 +0800
Subject: [PATCH 045/124] fix test.

---
 .../cpp/tree/gpu_hist/test_row_partitioner.cu | 20 +++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index 9b16cca5362d..9a6b122b7d62 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -2,6 +2,7 @@
  * Copyright 2019-2021 by XGBoost Contributors
  */
 #include <gtest/gtest.h>
+#include <cstddef>
 #include <vector>
 
 #include <thrust/device_vector.h>
@@ -10,6 +11,10 @@
 
 #include "../../../../src/tree/gpu_hist/row_partitioner.cuh"
 #include "../../helpers.h"
+#include "xgboost/base.h"
+#include "xgboost/generic_parameters.h"
+#include "xgboost/task.h"
+#include "xgboost/tree_model.h"
 
 namespace xgboost {
 namespace tree {
@@ -104,10 +109,17 @@ TEST(RowPartitioner, Basic) { TestUpdatePosition(); }
 void TestFinalise() {
   const int kNumRows = 10;
   RowPartitioner rp(0, kNumRows);
-  rp.FinalisePosition([=]__device__(RowPartitioner::RowIndexT ridx, int position)
-  {
-    return 7;
-  });
+
+  ObjInfo task{ObjInfo::kRegression, false, false};
+  RegTree tree;
+  std::vector<RowIndexCache> row_index;
+  Context ctx;
+
+  rp.FinalisePosition(
+      &ctx, &tree, task, &row_index,
+      [=] __device__(RowPartitioner::RowIndexT ridx, int position) { return 7; },
+      [] XGBOOST_DEVICE(size_t idx) { return false; });
+
   auto position = rp.GetPositionHost();
   for(auto p:position)
   {

From 10b210841981e918ef12ae9cbaf0d63cf8843ff7 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 10:37:01 +0800
Subject: [PATCH 046/124] Fix test.

---
 tests/cpp/tree/test_gpu_hist.cu | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/tests/cpp/tree/test_gpu_hist.cu b/tests/cpp/tree/test_gpu_hist.cu
index 82f40465deb2..62013290880d 100644
--- a/tests/cpp/tree/test_gpu_hist.cu
+++ b/tests/cpp/tree/test_gpu_hist.cu
@@ -13,6 +13,7 @@
 #include "../helpers.h"
 #include "../histogram_helpers.h"
 
+#include "xgboost/generic_parameters.h"
 #include "xgboost/json.h"
 #include "../../../src/data/sparse_page_source.h"
 #include "../../../src/tree/updater_gpu_hist.cu"
@@ -22,6 +23,13 @@
 
 namespace xgboost {
 namespace tree {
+namespace {
+auto MakeCtx() {
+  Context ctx;
+  ctx.gpu_id = 0;
+  return ctx;
+}
+}  // anonymous namespace
 
 TEST(GpuHist, DeviceHistogram) {
   // Ensures that node allocates correctly after reaching `kStopGrowingSize`.
@@ -81,8 +89,9 @@ void TestBuildHist(bool use_shared_memory_histograms) {
   param.Init(args);
   auto page = BuildEllpackPage(kNRows, kNCols);
   BatchParam batch_param{};
-  GPUHistMakerDevice<GradientSumT> maker(0, page.get(), {}, kNRows, param,
-                                         kNCols, kNCols, batch_param);
+  Context ctx{MakeCtx()};
+  GPUHistMakerDevice<GradientSumT> maker(&ctx, page.get(), {}, kNRows, param, kNCols, kNCols,
+                                         batch_param);
   xgboost::SimpleLCG gen;
   xgboost::SimpleRealUniformDistribution<bst_float> dist(0.0f, 1.0f);
   HostDeviceVector<GradientPair> gpair(kNRows);
@@ -158,14 +167,14 @@ TEST(GpuHist, ApplySplit) {
   BatchParam bparam;
   bparam.gpu_id = 0;
   bparam.max_bin = 3;
+  Context ctx{MakeCtx()};
 
   for (auto& ellpack : m->GetBatches<EllpackPage>(bparam)){
     auto impl = ellpack.Impl();
     HostDeviceVector<FeatureType> feature_types(10, FeatureType::kCategorical);
     feature_types.SetDevice(bparam.gpu_id);
     tree::GPUHistMakerDevice<GradientPairPrecise> updater(
-        0, impl, feature_types.ConstDeviceSpan(), n_rows, tparam, 0, n_cols,
-        bparam);
+        &ctx, impl, feature_types.ConstDeviceSpan(), n_rows, tparam, 0, n_cols, bparam);
     updater.ApplySplit(candidate, &tree);
 
     ASSERT_EQ(tree.GetSplitTypes().size(), 3);
@@ -224,8 +233,9 @@ TEST(GpuHist, EvaluateRootSplit) {
   // Initialize GPUHistMakerDevice
   auto page = BuildEllpackPage(kNRows, kNCols);
   BatchParam batch_param{};
-  GPUHistMakerDevice<GradientPairPrecise> maker(
-      0, page.get(), {}, kNRows, param, kNCols, kNCols, batch_param);
+  Context ctx{MakeCtx()};
+  GPUHistMakerDevice<GradientPairPrecise> maker(&ctx, page.get(), {}, kNRows, param, kNCols, kNCols,
+                                                batch_param);
   // Initialize GPUHistMakerDevice::node_sum_gradients
   maker.node_sum_gradients = {};
 

From aea2ead7cbb1415a4403e8bd9a076824485c718b Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 11:13:46 +0800
Subject: [PATCH 047/124] Start working on test.

---
 tests/cpp/common/test_stats.cu | 32 +++++++++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/tests/cpp/common/test_stats.cu b/tests/cpp/common/test_stats.cu
index 777c8b9058db..6aa983c7c481 100644
--- a/tests/cpp/common/test_stats.cu
+++ b/tests/cpp/common/test_stats.cu
@@ -5,17 +5,43 @@
 #include <cstddef>
 
 #include "../../../src/common/stats.cuh"
+#include "xgboost/base.h"
+#include "xgboost/generic_parameters.h"
 #include "xgboost/host_device_vector.h"
+#include "xgboost/linalg.h"
+
 namespace xgboost {
 namespace common {
-TEST(Stats, GPUQuantile) {
-  linalg::Tensor<float, 1> arr({20.f, 0.f, 15.f, 50.f, 40.f, 0.f, 35.f}, {7}, Context::kCpuId);
+namespace {
+void TestStatsGPUQuantile() {
+  linalg::Tensor<float, 1> arr(
+      {1.f, 2.f, 3.f, 4.f, 5.f,
+       2.f, 4.f, 5.f, 3.f, 1.f},
+      {10}, 0);
+  linalg::Tensor<size_t, 1> indptr({0, 5, 10}, {3}, 0);
   HostDeviceVector<float> resutls;
 
   auto d_arr = arr.View(0);
+  auto d_key = indptr.View(0);
+
+  auto key_it = dh::MakeTransformIterator<size_t>(
+      thrust::make_counting_iterator(0ul), [=] __device__(size_t i) { return d_key(i); });
   auto val_it = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
-                                                 [=](size_t i) { return d_arr(i); });
+                                                 [=] XGBOOST_DEVICE(size_t i) { return d_arr(i); });
+
+  Context ctx;
+  ctx.gpu_id = 0;
+  SegmentedQuantile(&ctx, 0.5, key_it, key_it + indptr.Size(), val_it, val_it + arr.Size(),
+                    &resutls);
+
+  auto const& h_results = resutls.HostVector();
+  ASSERT_EQ(h_results.size(), indptr.Size() - 1);
+  ASSERT_EQ(h_results.front(), 3.0f);
+  ASSERT_EQ(h_results.back(), 3.0f);
 }
+}  // anonymous namespace
+
+TEST(Stats, GPUQuantile) { TestStatsGPUQuantile(); }
 
 TEST(Stats, GPUWeightedQuantile) {}
 }  // namespace common

From 13b5f286a621fb479589b223952a46c7e72ec7f0 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 11:19:50 +0800
Subject: [PATCH 048/124] Avoid computing residule.

---
 src/objective/regression_obj.cu | 35 +++++++--------------------------
 1 file changed, 7 insertions(+), 28 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 34577b93dfdd..953a339d37ad 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -672,27 +672,6 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 
 #if defined(XGBOOST_USE_CUDA)
 namespace detail {
-/**
- * \brief Compute the residue between label and prediction.  Can be simplifed once we have
- *        prediction cache as matrix.
- */
-inline void ResidulesPredtY(Context const* ctx, linalg::TensorView<float const, 2> d_labels,
-                            common::Span<float const> d_predt,
-                            linalg::Tensor<float, 2>* p_residue) {
-  linalg::Tensor<float, 2>& residue = *p_residue;
-  residue.SetDevice(ctx->gpu_id);
-  residue.Reshape(d_labels.Shape(0), d_labels.Shape(1));
-
-  auto d_residue = residue.View(ctx->gpu_id);
-  CHECK_EQ(d_predt.size(), d_labels.Size());
-  linalg::ElementWiseKernel(ctx, d_labels, [=] XGBOOST_DEVICE(size_t i, float y) mutable {
-    auto idx = linalg::UnravelIndex(i, d_labels.Shape());
-    size_t sample_id = std::get<0>(idx);
-    size_t target_id = std::get<1>(idx);
-    d_residue(sample_id, target_id) = y - d_predt[i];
-  });
-}
-
 void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const> row_index,
                           MetaInfo const& info, HostDeviceVector<float> const& predt,
                           uint32_t target, float alpha, RegTree* p_tree) {
@@ -704,19 +683,19 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   HostDeviceVector<float> results;
   auto d_predt = predt.ConstDeviceSpan();
   auto d_labels = info.labels.View(ctx->gpu_id);
-  linalg::Tensor<float, 2> residue;
-  ResidulesPredtY(ctx, d_labels, d_predt, &residue);
-  auto d_residue = residue.View(ctx->gpu_id);
 
   auto d_row_index = part.row_index.ConstDeviceSpan();
   auto key_it = part.node_ptr.ConstDeviceSpan().data();
-  auto val_it = dh::MakeTransformIterator<float>(
-      thrust::make_counting_iterator(0ul),
-      [=] XGBOOST_DEVICE(size_t i) { return d_residue(d_row_index[i]); });
+  auto val_it = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
+                                                 [=] XGBOOST_DEVICE(size_t i) {
+                                                   auto predt = d_predt[d_row_index[i]];
+                                                   auto y = d_labels(d_row_index[i]);
+                                                   return y - predt;
+                                                 });
   CHECK_EQ(part.node_idx.Size() + 1, part.node_ptr.Size());
   if (info.weights_.Empty()) {
     common::SegmentedQuantile(ctx, alpha, key_it, key_it + part.node_ptr.Size(), val_it,
-                              val_it + d_residue.Size(), &results);
+                              val_it + d_labels.Size(), &results);
   } else {
     common::SegmentedWeightedQuantile(ctx, alpha, part, info, predt, &results);
   }

From 965060c8e33af01345f1079a5d86713bc69b8d36 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 12:05:52 +0800
Subject: [PATCH 049/124] Use iter for weighted.

---
 src/common/stats.cuh            | 130 +++++++-------------------------
 src/objective/regression_obj.cu |  30 ++++----
 2 files changed, 46 insertions(+), 114 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index fba291c7e6cd..5645d09ad522 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -7,6 +7,7 @@
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/sort.h>
 
+#include <cstddef>
 #include <iterator>  // std::distance
 
 #include "device_helpers.cuh"
@@ -18,67 +19,9 @@
 namespace xgboost {
 namespace common {
 namespace detail {
-/**
- * \brief Compute the residue between label and prediction.  Can be simplifed once we have
- *        prediction cache as matrix.
- */
-inline void ResidulesPredtY(Context const* ctx, linalg::TensorView<float const, 2> d_labels,
-                            common::Span<float const> d_predt,
-                            linalg::Tensor<float, 2>* p_residue) {
-  linalg::Tensor<float, 2>& residue = *p_residue;
-  residue.SetDevice(ctx->gpu_id);
-  residue.Reshape(d_labels.Shape(0), d_labels.Shape(1));
-
-  auto d_residue = residue.View(ctx->gpu_id);
-  CHECK_EQ(d_predt.size(), d_labels.Size());
-  linalg::ElementWiseKernel(ctx, d_labels, [=] XGBOOST_DEVICE(size_t i, float y) mutable {
-    auto idx = linalg::UnravelIndex(i, d_labels.Shape());
-    size_t sample_id = std::get<0>(idx);
-    size_t target_id = std::get<1>(idx);
-    d_residue(sample_id, target_id) = y - d_predt[i];
-  });
-}
-
-/**
- * \brief Argsort based on residue value and row partition.
- */
-template <typename Tup>
-inline void SortLeafRows(linalg::TensorView<float const, 2> d_residue,
-                         RowIndexCache const& row_index, dh::device_vector<size_t>* p_sorted_idx,
-                         dh::device_vector<Tup>* p_keys) {
-  auto& sorted_idx = *p_sorted_idx;
-  sorted_idx.resize(d_residue.Shape(0));
-  dh::Iota(dh::ToSpan(sorted_idx));
-
-  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
-  auto d_row_index = row_index.row_index.ConstDeviceSpan();
-  auto key_it = dh::MakeTransformIterator<Tup>(
-      thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
-        auto idx = linalg::UnravelIndex(i, d_residue.Shape());
-        size_t sample_id = std::get<0>(idx);
-        size_t target_id = std::get<1>(idx);
-        auto leaf_idx = dh::SegmentId(d_leaf_ptr, sample_id);
-        auto residue = d_residue(d_row_index[sample_id], target_id);
-        return thrust::make_tuple(leaf_idx, residue);
-      });
-  dh::device_vector<Tup>& keys = *p_keys;
-  keys.resize(d_residue.Size());
-  dh::XGBCachingDeviceAllocator<char> caching;
-  thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
-
-  dh::XGBDeviceAllocator<char> alloc;
-  thrust::stable_sort_by_key(thrust::cuda::par(alloc), keys.begin(), keys.end(), sorted_idx.begin(),
-                             [=] XGBOOST_DEVICE(Tup const& l, Tup const& r) {
-                               if (thrust::get<0>(l) != thrust::get<0>(r)) {
-                                 return thrust::get<0>(l) < thrust::get<0>(r);  // segment index
-                               }
-                               return thrust::get<1>(l) < thrust::get<1>(r);  // residue
-                             });
-}
-
-template <typename KeyIt, typename ValIt>
-inline void SortOnSegmented(KeyIt key_begin, KeyIt key_end, ValIt val_begin, ValIt val_end,
-                            dh::device_vector<size_t>* p_sorted_idx) {
+template <typename SegIt, typename ValIt>
+inline void SegmentedArgSort(SegIt seg_begin, SegIt seg_end, ValIt val_begin, ValIt val_end,
+                             dh::device_vector<size_t>* p_sorted_idx) {
   using Tup = thrust::tuple<size_t, float>;
   auto& sorted_idx = *p_sorted_idx;
   size_t n = std::distance(val_begin, val_end);
@@ -87,7 +30,7 @@ inline void SortOnSegmented(KeyIt key_begin, KeyIt key_end, ValIt val_begin, Val
   dh::device_vector<Tup> keys(sorted_idx.size());
   auto key_it = dh::MakeTransformIterator<Tup>(
       thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
-        auto leaf_idx = dh::SegmentId(key_begin, key_end, i);
+        auto leaf_idx = dh::SegmentId(seg_begin, seg_end, i);
         auto residue = val_begin[i];
         return thrust::make_tuple(leaf_idx, residue);
       });
@@ -108,22 +51,22 @@ inline void SortOnSegmented(KeyIt key_begin, KeyIt key_end, ValIt val_begin, Val
 /**
  * \brief Compute segmented quantile on GPU.
  *
- * \tparam PtrIt Iterator for CSR style segments indptr
+ * \tparam SegIt Iterator for CSR style segments indptr
  * \tparam ValIt Iterator for values
  *
  * \param alpha The p^th quantile we want to compute
  *
  *    std::distance(ptr_begin, ptr_end) should be equal to n_segments + 1
  */
-template <typename PtrIt, typename ValIt>
-void SegmentedQuantile(Context const* ctx, double alpha, PtrIt key_begin, PtrIt key_end,
+template <typename SegIt, typename ValIt>
+void SegmentedQuantile(Context const* ctx, double alpha, SegIt seg_begin, SegIt seg_end,
                        ValIt val_begin, ValIt val_end, HostDeviceVector<float>* quantiles) {
   CHECK(alpha >= 0 && alpha <= 1);
 
   dh::device_vector<size_t> sorted_idx;
   using Tup = thrust::tuple<size_t, float>;
-  detail::SortOnSegmented(key_begin, key_end, val_begin, val_end, &sorted_idx);
-  auto n_segments = std::distance(key_begin, key_end) - 1;
+  detail::SegmentedArgSort(seg_begin, seg_end, val_begin, val_end, &sorted_idx);
+  auto n_segments = std::distance(seg_begin, seg_end) - 1;
   if (n_segments <= 0) {
     return;
   }
@@ -136,15 +79,15 @@ void SegmentedQuantile(Context const* ctx, double alpha, PtrIt key_begin, PtrIt
   dh::LaunchN(n_segments, [=] XGBOOST_DEVICE(size_t i) {
     // each segment is the index of a leaf.
     size_t seg_idx = i;
-    size_t begin = key_begin[seg_idx];
-    auto n = static_cast<double>(key_begin[seg_idx + 1] - begin);
+    size_t begin = seg_begin[seg_idx];
+    auto n = static_cast<double>(seg_begin[seg_idx + 1] - begin);
 
     if (alpha <= (1 / (n + 1))) {
       d_results[i] = val_begin[d_sorted_idx[begin]];
       return;
     }
     if (alpha >= (n / (n + 1))) {
-      d_results[i] = val_begin[common::LastOf(seg_idx, key_begin)];
+      d_results[i] = val_begin[common::LastOf(seg_idx, seg_begin)];
       return;
     }
 
@@ -157,52 +100,37 @@ void SegmentedQuantile(Context const* ctx, double alpha, PtrIt key_begin, PtrIt
   });
 }
 
-inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
-                                      RowIndexCache const& row_index, MetaInfo const& info,
-                                      HostDeviceVector<float> const& predt,
-                                      HostDeviceVector<float>* quantiles) {
+template <typename SegIt, typename ValIt, typename WIter>
+void SegmentedWeightedQuantile(Context const* ctx, double alpha, SegIt seg_beg, SegIt seg_end,
+                               ValIt val_begin, ValIt val_end, WIter w_begin, WIter w_end,
+                               HostDeviceVector<float>* quantiles) {
   CHECK(alpha >= 0 && alpha <= 1);
-  auto d_predt = predt.ConstDeviceSpan();
-  auto d_labels = info.labels.View(ctx->gpu_id);
-  linalg::Tensor<float, 2> residue{d_labels.Shape(), ctx->gpu_id};
-  detail::ResidulesPredtY(ctx, d_labels, d_predt, &residue);
-  auto d_residue = residue.View(ctx->gpu_id);
-
   dh::device_vector<size_t> sorted_idx;
-  using Tup = thrust::tuple<size_t, float>;
-  dh::device_vector<Tup> keys;
-  detail::SortLeafRows(d_residue, row_index, &sorted_idx, &keys);
+  detail::SegmentedArgSort(seg_beg, seg_end, val_begin, val_end, &sorted_idx);
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
-  auto d_row_index = row_index.row_index.ConstDeviceSpan();
-  auto d_leaf_ptr = row_index.node_ptr.ConstDeviceSpan();
-  auto d_keys = dh::ToSpan(keys);
-
-  auto weights = info.weights_.ConstDeviceSpan();
-
-  dh::device_vector<float> weights_cdf(weights.size());
-  CHECK_EQ(weights.size(), d_labels.Shape(0));
+  size_t n_samples = std::distance(w_begin, w_end);
+  dh::device_vector<float> weights_cdf(n_samples);
 
   dh::XGBCachingDeviceAllocator<char> caching;
-  Span<size_t const> node_ptr = row_index.node_ptr.ConstDeviceSpan();
   auto scan_key = dh::MakeTransformIterator<size_t>(
       thrust::make_counting_iterator(0ul),
-      [=] XGBOOST_DEVICE(size_t i) { return dh::SegmentId(d_leaf_ptr, i); });
+      [=] XGBOOST_DEVICE(size_t i) { return dh::SegmentId(seg_beg, seg_end, i); });
   auto scan_val = dh::MakeTransformIterator<float>(
       thrust::make_counting_iterator(0ul),
-      [=] XGBOOST_DEVICE(size_t i) { return weights[d_row_index[d_sorted_idx[i]]]; });
-  thrust::inclusive_scan_by_key(thrust::cuda::par(caching), scan_key, scan_key + weights.size(),
+      [=] XGBOOST_DEVICE(size_t i) { return w_begin[d_sorted_idx[i]]; });
+  thrust::inclusive_scan_by_key(thrust::cuda::par(caching), scan_key, scan_key + n_samples,
                                 scan_val, weights_cdf.begin());
 
+  auto n_segments = std::distance(seg_beg, seg_end) - 1;
   quantiles->SetDevice(ctx->gpu_id);
-  quantiles->Resize(row_index.node_idx.Size());
+  quantiles->Resize(n_segments);
   auto d_results = quantiles->DeviceSpan();
   auto d_weight_cdf = dh::ToSpan(weights_cdf);
 
-  dh::LaunchN(row_index.node_idx.Size(), [=] XGBOOST_DEVICE(size_t i) {
-    size_t target_id = 0;
+  dh::LaunchN(n_segments, [=] XGBOOST_DEVICE(size_t i) {
     size_t seg_idx = i;
-    size_t begin = d_leaf_ptr[seg_idx];
-    auto n = static_cast<double>(d_leaf_ptr[seg_idx + 1] - begin);
+    size_t begin = seg_beg[seg_idx];
+    auto n = static_cast<double>(seg_beg[seg_idx + 1] - begin);
     auto leaf_cdf = d_weight_cdf.subspan(begin, static_cast<size_t>(n));
     auto leaf_sorted_idx = d_sorted_idx.subspan(begin, static_cast<size_t>(n));
     float thresh = leaf_cdf.back() * alpha;
@@ -211,7 +139,7 @@ inline void SegmentedWeightedQuantile(Context const* ctx, double alpha,
                                      leaf_cdf.data() + leaf_cdf.size(), thresh) -
                  leaf_cdf.data();
     idx = std::min(idx, static_cast<size_t>(n - 1));
-    d_results[i] = d_residue(d_row_index[leaf_sorted_idx[idx]], target_id);
+    d_results[i] = val_begin[leaf_sorted_idx[idx]];
   });
 }
 }  // namespace common
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 953a339d37ad..44782e9ec4f9 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -685,19 +685,24 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   auto d_labels = info.labels.View(ctx->gpu_id);
 
   auto d_row_index = part.row_index.ConstDeviceSpan();
-  auto key_it = part.node_ptr.ConstDeviceSpan().data();
-  auto val_it = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
-                                                 [=] XGBOOST_DEVICE(size_t i) {
-                                                   auto predt = d_predt[d_row_index[i]];
-                                                   auto y = d_labels(d_row_index[i]);
-                                                   return y - predt;
-                                                 });
+  auto seg_beg = part.node_ptr.ConstDeviceSpan().data();
+  auto seg_end = seg_beg + part.node_ptr.Size();
+  auto val_beg = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
+                                                  [=] XGBOOST_DEVICE(size_t i) {
+                                                    auto predt = d_predt[d_row_index[i]];
+                                                    auto y = d_labels(d_row_index[i]);
+                                                    return y - predt;
+                                                  });
+  auto val_end = val_beg + d_labels.Size();
   CHECK_EQ(part.node_idx.Size() + 1, part.node_ptr.Size());
   if (info.weights_.Empty()) {
-    common::SegmentedQuantile(ctx, alpha, key_it, key_it + part.node_ptr.Size(), val_it,
-                              val_it + d_labels.Size(), &results);
+    common::SegmentedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, &results);
   } else {
-    common::SegmentedWeightedQuantile(ctx, alpha, part, info, predt, &results);
+    auto d_weights = info.weights_.ConstDeviceSpan();
+    CHECK_EQ(d_weights.size(), d_row_index.size());
+    auto w_it = thrust::make_permutation_iterator(dh::tcbegin(d_weights), dh::tcbegin(d_row_index));
+    common::SegmentedWeightedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, w_it,
+                                      w_it + d_weights.size(), &results);
   }
 
   auto const& h_results = results.HostVector();
@@ -779,12 +784,11 @@ class MeanAbsoluteError : public ObjFunction {
     return std::max(static_cast<size_t>(1), info.labels.Shape(1));
   }
 
-  ObjInfo Task() const override {
-    return {ObjInfo::kRegression, true, true};
-  }
+  ObjInfo Task() const override { return {ObjInfo::kRegression, true, true}; }
 
   void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
                    HostDeviceVector<GradientPair>* out_gpair) override {
+    CheckRegInputs(info, preds);
     auto labels = info.labels.View(ctx_->gpu_id);
 
     out_gpair->SetDevice(ctx_->gpu_id);

From b6257bd72cc16911ff3bad9d1b5d33bf8fd465e7 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 12:37:36 +0800
Subject: [PATCH 050/124] Use permutation iterator.

---
 src/common/stats.cuh           | 11 +++--
 tests/cpp/common/test_stats.cu | 84 +++++++++++++++++++++++-----------
 2 files changed, 64 insertions(+), 31 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index 5645d09ad522..9311a82cb3df 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -5,6 +5,7 @@
 #define XGBOOST_COMMON_STATS_CUH_
 
 #include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/permutation_iterator.h>
 #include <thrust/sort.h>
 
 #include <cstddef>
@@ -76,6 +77,8 @@ void SegmentedQuantile(Context const* ctx, double alpha, SegIt seg_begin, SegIt
   auto d_results = quantiles->DeviceSpan();
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
 
+  auto val = thrust::make_permutation_iterator(val_begin, dh::tcbegin(d_sorted_idx));
+
   dh::LaunchN(n_segments, [=] XGBOOST_DEVICE(size_t i) {
     // each segment is the index of a leaf.
     size_t seg_idx = i;
@@ -83,19 +86,19 @@ void SegmentedQuantile(Context const* ctx, double alpha, SegIt seg_begin, SegIt
     auto n = static_cast<double>(seg_begin[seg_idx + 1] - begin);
 
     if (alpha <= (1 / (n + 1))) {
-      d_results[i] = val_begin[d_sorted_idx[begin]];
+      d_results[i] = val[begin];
       return;
     }
     if (alpha >= (n / (n + 1))) {
-      d_results[i] = val_begin[common::LastOf(seg_idx, seg_begin)];
+      d_results[i] = val[common::LastOf(seg_idx, seg_begin)];
       return;
     }
 
     double x = alpha * static_cast<double>(n + 1);
     double k = std::floor(x) - 1;
     double d = (x - 1) - k;
-    auto v0 = val_begin[d_sorted_idx[begin + static_cast<size_t>(k)]];
-    auto v1 = val_begin[d_sorted_idx[begin + static_cast<size_t>(k) + 1]];
+    auto v0 = val[begin + static_cast<size_t>(k)];
+    auto v1 = val[begin + static_cast<size_t>(k) + 1];
     d_results[seg_idx] = v0 + d * (v1 - v0);
   });
 }
diff --git a/tests/cpp/common/test_stats.cu b/tests/cpp/common/test_stats.cu
index 6aa983c7c481..fb28d8e4da98 100644
--- a/tests/cpp/common/test_stats.cu
+++ b/tests/cpp/common/test_stats.cu
@@ -3,6 +3,8 @@
  */
 #include <gtest/gtest.h>
 #include <cstddef>
+#include <utility>
+#include <vector>
 
 #include "../../../src/common/stats.cuh"
 #include "xgboost/base.h"
@@ -13,36 +15,64 @@
 namespace xgboost {
 namespace common {
 namespace {
-void TestStatsGPUQuantile() {
-  linalg::Tensor<float, 1> arr(
+class StatsGPU : public ::testing::Test {
+ private:
+  linalg::Tensor<float, 1> arr_{
       {1.f, 2.f, 3.f, 4.f, 5.f,
        2.f, 4.f, 5.f, 3.f, 1.f},
-      {10}, 0);
-  linalg::Tensor<size_t, 1> indptr({0, 5, 10}, {3}, 0);
-  HostDeviceVector<float> resutls;
-
-  auto d_arr = arr.View(0);
-  auto d_key = indptr.View(0);
-
-  auto key_it = dh::MakeTransformIterator<size_t>(
-      thrust::make_counting_iterator(0ul), [=] __device__(size_t i) { return d_key(i); });
-  auto val_it = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
-                                                 [=] XGBOOST_DEVICE(size_t i) { return d_arr(i); });
-
-  Context ctx;
-  ctx.gpu_id = 0;
-  SegmentedQuantile(&ctx, 0.5, key_it, key_it + indptr.Size(), val_it, val_it + arr.Size(),
-                    &resutls);
-
-  auto const& h_results = resutls.HostVector();
-  ASSERT_EQ(h_results.size(), indptr.Size() - 1);
-  ASSERT_EQ(h_results.front(), 3.0f);
-  ASSERT_EQ(h_results.back(), 3.0f);
-}
-}  // anonymous namespace
+      {10}, 0};
+  linalg::Tensor<size_t, 1> indptr_{{0, 5, 10}, {3}, 0};
+  HostDeviceVector<float> resutls_;
+  using TestSet = std::vector<std::pair<float, float>>;
+  Context ctx_;
+
+  void Check(float expected) {
+    auto const& h_results = resutls_.HostVector();
+    ASSERT_EQ(h_results.size(), indptr_.Size() - 1);
+    ASSERT_EQ(h_results.front(), expected);
+    EXPECT_EQ(h_results.back(), expected);
+  }
+
+ public:
+  void SetUp() override { ctx_.gpu_id = 0; }
+  void Weighted() {
+    auto d_arr = arr_.View(0);
+    auto d_key = indptr_.View(0);
+
+    auto key_it = dh::MakeTransformIterator<size_t>(thrust::make_counting_iterator(0ul),
+                                                    [=] __device__(size_t i) { return d_key(i); });
+    auto val_it = dh::MakeTransformIterator<float>(
+        thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) { return d_arr(i); });
+    linalg::Tensor<float, 1> weights{{10}, 0};
+    linalg::ElementWiseTransformDevice(weights.View(0),
+                                       [=] XGBOOST_DEVICE(size_t, float) { return 1.0; });
+    auto w_it = weights.Data()->ConstDevicePointer();
+    for (auto const& pair : TestSet{{0.0f, 1.0f}, {0.5f, 3.0f}, {1.0f, 5.0f}}) {
+      SegmentedWeightedQuantile(&ctx_, pair.first, key_it, key_it + indptr_.Size(), val_it,
+                                val_it + arr_.Size(), w_it, w_it + weights.Size(), &resutls_);
+      this->Check(pair.second);
+    }
+  }
 
-TEST(Stats, GPUQuantile) { TestStatsGPUQuantile(); }
+  void NonWeighted() {
+    auto d_arr = arr_.View(0);
+    auto d_key = indptr_.View(0);
+
+    auto key_it = dh::MakeTransformIterator<size_t>(thrust::make_counting_iterator(0ul),
+                                                    [=] __device__(size_t i) { return d_key(i); });
+    auto val_it = dh::MakeTransformIterator<float>(
+        thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) { return d_arr(i); });
+
+    for (auto const& pair : TestSet{{0.0f, 1.0f}, {0.5f, 3.0f}, {1.0f, 5.0f}}) {
+      SegmentedQuantile(&ctx_, pair.first, key_it, key_it + indptr_.Size(), val_it,
+                        val_it + arr_.Size(), &resutls_);
+      this->Check(pair.second);
+    }
+  }
+};
+}  // anonymous namespace
 
-TEST(Stats, GPUWeightedQuantile) {}
+TEST_F(StatsGPU, Quantile) { this->NonWeighted(); }
+TEST_F(StatsGPU, WeightedQuantile) { this->Weighted(); }
 }  // namespace common
 }  // namespace xgboost

From 321558f9aec748f41b7492c626207a6fc5935ada Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 12:43:11 +0800
Subject: [PATCH 051/124] Allreduce.

---
 src/objective/regression_obj.cu | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 44782e9ec4f9..eaf7b0d41481 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -705,12 +705,18 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
                                       w_it + d_weights.size(), &results);
   }
 
-  auto const& h_results = results.HostVector();
+  auto& quantiles = results.HostVector();
+  // FIXME(jiamingy): Use nccl once we have an unified allreducer.
+  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
+  auto world = rabit::GetWorldSize();
+  std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
+                 [&](float q) { return q / world; });
+
   auto& tree = *p_tree;
   auto const& h_node_idx = row_index.front().node_idx.HostVector();
   for (size_t i = 0; i < h_node_idx.size(); ++i) {
     auto nidx = h_node_idx[i];
-    auto q = h_results[i];
+    auto q = quantiles[i];
     CHECK(tree[nidx].IsLeaf());
     tree[nidx].SetLeaf(q);  // fixme: exact tree method
   }

From 05a8133c3ee6fd1cb0e07b526f415590775935d9 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 12:46:24 +0800
Subject: [PATCH 052/124] Disable multi-target.

---
 src/common/stats.h                    |  6 ++---
 src/objective/regression_obj.cu       | 17 +++++--------
 src/tree/gpu_hist/row_partitioner.cuh |  2 +-
 tests/cpp/common/test_stats.cc        | 36 ++++++++++++++-------------
 4 files changed, 29 insertions(+), 32 deletions(-)

diff --git a/src/common/stats.h b/src/common/stats.h
index de081a5b476e..5915cc867880 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -26,7 +26,7 @@ namespace common {
  * \return The result of interpolation.
  */
 template <typename Iter>
-float Percentile(double alpha, Iter const& begin, Iter const& end) {
+float Quantile(double alpha, Iter const& begin, Iter const& end) {
   CHECK(alpha >= 0 && alpha <= 1);
   auto n = static_cast<double>(std::distance(begin, end));
   if (n == 0) {
@@ -66,7 +66,7 @@ float Percentile(double alpha, Iter const& begin, Iter const& end) {
  *   weighted quantile with interpolation.
  */
 template <typename Iter, typename WeightIter>
-float WeightedPercentile(double quantile, Iter begin, Iter end, WeightIter weights) {
+float WeightedQuantile(double alpha, Iter begin, Iter end, WeightIter weights) {
   auto n = static_cast<double>(std::distance(begin, end));
   std::vector<size_t> sorted_idx(n);
   std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
@@ -81,7 +81,7 @@ float WeightedPercentile(double quantile, Iter begin, Iter end, WeightIter weigh
   for (size_t i = 1; i < n; ++i) {
     weight_cdf[i] = weight_cdf[i - 1] + *(weights + sorted_idx[i]);
   }
-  float thresh = weight_cdf.back() * quantile;
+  float thresh = weight_cdf.back() * alpha;
   size_t idx =
       std::lower_bound(weight_cdf.cbegin(), weight_cdf.cend(), thresh) - weight_cdf.cbegin();
   idx = std::min(idx, static_cast<size_t>(n - 1));
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index eaf7b0d41481..8bf33bd77c3e 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -725,8 +725,8 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
 #endif  // defined(XGBOOST_USE_CUDA)
 
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
-                        MetaInfo const& info, HostDeviceVector<float> const& prediction,
-                        uint32_t target, float alpha, RegTree* p_tree) {
+                        MetaInfo const& info, HostDeviceVector<float> const& predt, uint32_t target,
+                        float alpha, RegTree* p_tree) {
   auto& tree = *p_tree;
   std::vector<float> quantiles;
   for (auto const& part : row_index) {
@@ -736,12 +736,12 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
       CHECK(tree[seg.nidx].IsLeaf());
       auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
       auto h_labels = info.labels.HostView().Slice(linalg::All(), target);
-      auto const& h_prediction = prediction.ConstHostVector();
+      auto const& h_predt = predt.ConstHostVector();
       auto h_weights = linalg::MakeVec(&info.weights_);
 
       auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
         auto row_idx = h_row_set[i];
-        return h_labels(row_idx) - h_prediction[row_idx];
+        return h_labels(row_idx) - h_predt[row_idx];
       });
       auto w_it = common::MakeIndexTransformIter([&](size_t i) -> float {
         auto row_idx = h_row_set[i];
@@ -750,9 +750,9 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
 
       float q{0};
       if (info.weights_.Empty()) {
-        q = common::Percentile(alpha, iter, iter + h_row_set.size());
+        q = common::Quantile(alpha, iter, iter + h_row_set.size());
       } else {
-        q = common::WeightedPercentile(alpha, iter, iter + h_row_set.size(), w_it);
+        q = common::WeightedQuantile(alpha, iter, iter + h_row_set.size(), w_it);
       }
       results.at(k) = q;
     });
@@ -785,11 +785,6 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
 class MeanAbsoluteError : public ObjFunction {
  public:
   void Configure(Args const&) override {}
-
-  uint32_t Targets(MetaInfo const& info) const override {
-    return std::max(static_cast<size_t>(1), info.labels.Shape(1));
-  }
-
   ObjInfo Task() const override { return {ObjInfo::kRegression, true, true}; }
 
   void GetGradient(HostDeviceVector<bst_float> const& preds, const MetaInfo& info, int iter,
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 92f5b6cd04ca..40c52321cf75 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -43,7 +43,6 @@ class RowPartitioner {
 
  private:
   int device_idx_;
-  std::vector<Segment> ridx_segments_;
   /*! \brief In here if you want to find the rows belong to a node nid, first you need to
    * get the indices segment from ridx_segments[nid], then get the row index that
    * represents position of row in input data X.  `RowPartitioner::GetRows` would be a
@@ -52,6 +51,7 @@ class RowPartitioner {
    * node id -> segment -> indices of rows belonging to node
    */
   /*! \brief Range of row index for each node, pointers into ridx below. */
+  std::vector<Segment> ridx_segments_;
   dh::TemporaryArray<RowIndexT> ridx_a_;
   dh::TemporaryArray<RowIndexT> ridx_b_;
   dh::TemporaryArray<bst_node_t> position_a_;
diff --git a/tests/cpp/common/test_stats.cc b/tests/cpp/common/test_stats.cc
index ee446ea6a049..121dce4cba09 100644
--- a/tests/cpp/common/test_stats.cc
+++ b/tests/cpp/common/test_stats.cc
@@ -9,26 +9,28 @@
 namespace xgboost {
 namespace common {
 TEST(Stats, Percentil) {
-  linalg::Tensor<float, 1> arr({20.f, 0.f, 15.f, 50.f, 40.f, 0.f, 35.f}, {7}, Context::kCpuId);
-  std::vector<size_t> index{0, 2, 3, 4, 6};
-  auto h_arr = arr.HostView();
-  auto beg = MakeIndexTransformIter([&](size_t i) { return h_arr(index[i]); });
-  auto end = beg + index.size();
-  auto percentile = Percentile(0.40f, beg, end);
-  ASSERT_EQ(percentile, 26.0);
+  {
+    linalg::Tensor<float, 1> arr({20.f, 0.f, 15.f, 50.f, 40.f, 0.f, 35.f}, {7}, Context::kCpuId);
+    std::vector<size_t> index{0, 2, 3, 4, 6};
+    auto h_arr = arr.HostView();
+    auto beg = MakeIndexTransformIter([&](size_t i) { return h_arr(index[i]); });
+    auto end = beg + index.size();
+    auto q = Quantile(0.40f, beg, end);
+    ASSERT_EQ(q, 26.0);
 
-  percentile = Percentile(0.20f, beg, end);
-  ASSERT_EQ(percentile, 16.0);
+    q = Quantile(0.20f, beg, end);
+    ASSERT_EQ(q, 16.0);
 
-  percentile = Percentile(0.10f, beg, end);
-  ASSERT_EQ(percentile, 15.0);
+    q = Quantile(0.10f, beg, end);
+    ASSERT_EQ(q, 15.0);
+  }
 
   {
     std::vector<float> vec{1., 2., 3., 4., 5.};
     auto beg = MakeIndexTransformIter([&](size_t i) { return vec[i]; });
-    auto end = beg + index.size();
-    auto percentile = Percentile(0.5f, beg, end);
-    ASSERT_EQ(percentile, 3.);
+    auto end = beg + vec.size();
+    auto q = Quantile(0.5f, beg, end);
+    ASSERT_EQ(q, 3.);
   }
 }
 
@@ -43,13 +45,13 @@ TEST(Stats, WeightedQuantile) {
   auto end = beg + arr.Size();
   auto w = MakeIndexTransformIter([&](size_t i) { return h_weight(i); });
 
-  auto q = WeightedPercentile(0.50f, beg, end, w);
+  auto q = WeightedQuantile(0.50f, beg, end, w);
   ASSERT_EQ(q, 3);
 
-  q = WeightedPercentile(0.0, beg, end, w);
+  q = WeightedQuantile(0.0, beg, end, w);
   ASSERT_EQ(q, 1);
 
-  q = WeightedPercentile(1.0, beg, end, w);
+  q = WeightedQuantile(1.0, beg, end, w);
   ASSERT_EQ(q, 5);
 }
 }  // namespace common

From f1670af6a408bb908151699cc10fe2287ee81dd2 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 13:21:09 +0800
Subject: [PATCH 053/124] Fix regen.

---
 src/objective/regression_loss.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/objective/regression_loss.h b/src/objective/regression_loss.h
index 8d9d661f23ef..cbc161e0840d 100644
--- a/src/objective/regression_loss.h
+++ b/src/objective/regression_loss.h
@@ -38,7 +38,7 @@ struct LinearSquareLoss {
   static const char* DefaultEvalMetric() { return "rmse"; }
 
   static const char* Name() { return "reg:squarederror"; }
-  static ObjInfo Info() { return ObjInfo::kRegression; }
+  static ObjInfo Info() { return {ObjInfo::kRegression, true, false}; }
 };
 
 struct SquaredLogError {

From 9cf36fedeb9e335502aa64b99491d2ef8056df61 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 14:42:48 +0800
Subject: [PATCH 054/124] Restore ext handling.

---
 src/tree/gpu_hist/row_partitioner.cuh | 43 ++++++++++++++++-----------
 src/tree/updater_gpu_hist.cu          | 25 +++++++++++++---
 tests/cpp/tree/test_gpu_hist.cu       |  2 +-
 3 files changed, 48 insertions(+), 22 deletions(-)

diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 40c52321cf75..c4e564660b7e 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -170,25 +170,34 @@ class RowPartitioner {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
     auto sorted_position = position_.Other();
-    dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
-      auto position = d_position[idx];
-      RowIndexT ridx = d_ridx[idx];
-      bst_node_t new_position = op(ridx, position);
-      if (sampledp(ridx)) {
-        // push to the end
-        sorted_position[ridx] = std::numeric_limits<bst_node_t>::max();
-      } else {
-        sorted_position[ridx] = new_position;
-      }
-
-      if (new_position == kIgnoredTreePosition) {
-        return;
-      }
-      d_position[idx] = new_position;
-    });
-
     if (!task.zero_hess) {
+      dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
+        auto position = d_position[idx];
+        RowIndexT ridx = d_ridx[idx];
+        bst_node_t new_position = op(ridx, position);
+        if (new_position == kIgnoredTreePosition) {
+          return;
+        }
+        d_position[idx] = new_position;
+      });
       return;
+    } else {
+      dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
+        auto position = d_position[idx];
+        RowIndexT ridx = d_ridx[idx];
+        bst_node_t new_position = op(ridx, position);
+        if (sampledp(ridx)) {
+          // push to the end
+          sorted_position[ridx] = std::numeric_limits<bst_node_t>::max();
+        } else {
+          sorted_position[ridx] = new_position;
+        }
+
+        if (new_position == kIgnoredTreePosition) {
+          return;
+        }
+        d_position[idx] = new_position;
+      });
     }
     // copy position to buffer
     size_t n_samples = position_.Size();
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 15cd3670305f..48b3f33b8c38 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -13,6 +13,7 @@
 #include <vector>
 
 #include "xgboost/base.h"
+#include "xgboost/data.h"
 #include "xgboost/generic_parameters.h"
 #include "xgboost/host_device_vector.h"
 #include "xgboost/parameter.h"
@@ -413,9 +414,21 @@ struct GPUHistMakerDevice {
 
     CHECK(p_out_row_indices->empty());
     p_out_row_indices->push_back(RowIndexCache{ctx_, p_fmat->Info().num_row_});
-    FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
-                           dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
-                           p_out_row_indices);
+    if (row_partitioner->GetRows().size() != p_fmat->Info().num_row_) {
+      row_partitioner.reset();  // Release the device memory first before reallocating
+      row_partitioner.reset(new RowPartitioner(ctx_->gpu_id, p_fmat->Info().num_row_));
+    }
+    if (page->n_rows == p_fmat->Info().num_row_) {
+      FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
+                             dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
+                             p_out_row_indices);
+    } else {
+      for (auto const& batch : p_fmat->GetBatches<EllpackPage>(batch_param)) {
+        FinalisePositionInPage(batch.Impl(), p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
+                               dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
+                               p_out_row_indices);
+      }
+    }
   }
 
   void FinalisePositionInPage(EllpackPageImpl const *page,
@@ -464,7 +477,11 @@ struct GPUHistMakerDevice {
 
           return position;
         },
-        [d_gpair] __device__(size_t ridx) { return d_gpair[ridx].GetHess() - .0f == 0.f; });
+        [d_gpair] __device__(size_t ridx) {
+          // FIXME(jiamingy): Doesn't work when sampling is used with external memory as
+          // the sampler compacts the gradient vector.
+          return d_gpair[ridx].GetHess() - .0f == 0.f;
+        });
   }
 
   void UpdatePredictionCache(linalg::VectorView<float> out_preds_d, RegTree const* p_tree) {
diff --git a/tests/cpp/tree/test_gpu_hist.cu b/tests/cpp/tree/test_gpu_hist.cu
index 62013290880d..20a1978b08bf 100644
--- a/tests/cpp/tree/test_gpu_hist.cu
+++ b/tests/cpp/tree/test_gpu_hist.cu
@@ -493,7 +493,7 @@ TEST(GpuHist, ExternalMemoryWithSampling) {
   auto preds_h = preds.ConstHostVector();
   auto preds_ext_h = preds_ext.ConstHostVector();
   for (int i = 0; i < kRows; i++) {
-    EXPECT_NEAR(preds_h[i], preds_ext_h[i], 1e-3);
+    ASSERT_NEAR(preds_h[i], preds_ext_h[i], 1e-3);
   }
 }
 

From 3fa3bf88451e4f45910aae0676ec307a955ed443 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 14:52:23 +0800
Subject: [PATCH 055/124] Guard.

---
 include/xgboost/task.h       | 4 ++++
 src/gbm/gbtree.cc            | 2 +-
 src/tree/updater_gpu_hist.cu | 4 ++++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/include/xgboost/task.h b/include/xgboost/task.h
index 91b02a52e438..22a65f2dad4b 100644
--- a/include/xgboost/task.h
+++ b/include/xgboost/task.h
@@ -42,6 +42,10 @@ struct ObjInfo {
   XGBOOST_DEVICE bool UseOneHot() const {
     return (task != ObjInfo::kRegression && task != ObjInfo::kBinary);
   }
+  /**
+   * \brief Use adaptive tree if the objective doesn't have valid hessian value.
+   */
+  XGBOOST_DEVICE bool UpdateTreeLeaf() const { return zero_hess; }
 };
 }  // namespace xgboost
 #endif  // XGBOOST_TASK_H_
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index f23527c1c834..2a0371a3ce09 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -220,7 +220,7 @@ void CopyGradient(HostDeviceVector<GradientPair> const* in_gpair, int32_t n_thre
 void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const& predictions,
                             ObjFunction const* obj, size_t gidx,
                             std::vector<std::unique_ptr<RegTree>>* p_trees) {
-  if (!obj || !obj->Task().zero_hess) {
+  if (!obj || !obj->Task().UpdateTreeLeaf()) {
     return;
   }
   auto& trees = *p_trees;
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 48b3f33b8c38..e8e5902838bb 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -418,6 +418,10 @@ struct GPUHistMakerDevice {
       row_partitioner.reset();  // Release the device memory first before reallocating
       row_partitioner.reset(new RowPartitioner(ctx_->gpu_id, p_fmat->Info().num_row_));
     }
+    if (task.UpdateTreeLeaf() && !p_fmat->SingleColBlock() && param.subsample != 1.0) {
+      // see comment in the `FinalisePositionInPage`.
+      LOG(FATAL) << "Current objective function can not be used with external memory";
+    }
     if (page->n_rows == p_fmat->Info().num_row_) {
       FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
                              dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,

From df562e2dd8f347ff58daa0c70bbd35fa0b835a13 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 16:29:05 +0800
Subject: [PATCH 056/124] Empty matrix test.

---
 tests/python/test_with_dask.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index 4e80409d4764..fd1fcebdf3e9 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -666,7 +666,8 @@ def test_empty_dmatrix_training_continuation(client: "Client") -> None:
 def run_empty_dmatrix_reg(client: "Client", parameters: dict) -> None:
     def _check_outputs(out: xgb.dask.TrainReturnT, predictions: np.ndarray) -> None:
         assert isinstance(out['booster'], xgb.dask.Booster)
-        assert len(out['history']['validation']['rmse']) == 2
+        for _, v in out['history']['validation'].items():
+            assert len(v) == 2
         assert isinstance(predictions, np.ndarray)
         assert predictions.shape[0] == 1
 
@@ -867,6 +868,8 @@ def test_empty_dmatrix(tree_method) -> None:
             parameters = {'tree_method': tree_method}
             run_empty_dmatrix_reg(client, parameters)
             run_empty_dmatrix_cls(client, parameters)
+            parameters = {'tree_method': tree_method, "objective": "reg:absoluteerror"}
+            run_empty_dmatrix_reg(client, parameters)
 
 
 async def run_from_dask_array_asyncio(scheduler_address: str) -> xgb.dask.TrainReturnT:

From 0f486c04f3cae362e79f959a8cd5588ba4c63019 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 17:12:17 +0800
Subject: [PATCH 057/124] Fix empty dmatrix on CPU.

---
 src/objective/regression_obj.cu | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 8bf33bd77c3e..0658a1bc0672 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -670,8 +670,9 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .describe("Tweedie regression for insurance data.")
 .set_body([]() { return new TweedieRegression(); });
 
-#if defined(XGBOOST_USE_CUDA)
+
 namespace detail {
+#if defined(XGBOOST_USE_CUDA)
 void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const> row_index,
                           MetaInfo const& info, HostDeviceVector<float> const& predt,
                           uint32_t target, float alpha, RegTree* p_tree) {
@@ -721,14 +722,14 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
     tree[nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }
-}  // namespace detail
 #endif  // defined(XGBOOST_USE_CUDA)
 
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
                         MetaInfo const& info, HostDeviceVector<float> const& predt, uint32_t target,
                         float alpha, RegTree* p_tree) {
   auto& tree = *p_tree;
-  std::vector<float> quantiles;
+  CHECK(!row_index.empty());
+  std::vector<float> quantiles(row_index.front().indptr.size(), 0);
   for (auto const& part : row_index) {
     std::vector<float> results(part.indptr.size());
     common::ParallelFor(part.indptr.size(), ctx->Threads(), [&](size_t k) {
@@ -758,14 +759,17 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     });
 
     // fixme: verify this is correct for external memory
-    if (quantiles.empty()) {
-      quantiles.resize(results.size(), 0);
-    }
     for (size_t i = 0; i < results.size(); ++i) {
       quantiles[i] += results[i];
     }
   }
 
+  size_t n_leaf{quantiles.size()};
+  rabit::Allreduce<rabit::op::Max>(&n_leaf, 1);
+  CHECK(quantiles.empty() || quantiles.size() == n_leaf);
+  if (quantiles.empty()) {
+    quantiles.resize(n_leaf);
+  }
   // use the mean value
   rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
   auto world = rabit::GetWorldSize();
@@ -777,10 +781,10 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     auto seg = row_index.front().indptr[i];
     auto q = quantiles[i];
     CHECK(tree[seg.nidx].IsLeaf());
-
     tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }
+}  // namespace detail
 
 class MeanAbsoluteError : public ObjFunction {
  public:
@@ -817,7 +821,7 @@ class MeanAbsoluteError : public ObjFunction {
                       HostDeviceVector<float> const& prediction, uint32_t target,
                       RegTree* p_tree) const override {
     if (ctx_->IsCPU()) {
-      UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
+      detail::UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
     } else {
 #if defined(XGBOOST_USE_CUDA)
       detail::UpdateTreeLeafDevice(ctx_, row_index, info, prediction, target, 0.5, p_tree);

From 4d77a9017d819278647bc194e899b9d8247bb251 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 18:23:26 +0800
Subject: [PATCH 058/124] Unify the cache.

---
 include/xgboost/tree_model.h    |  7 -------
 src/common/partition_builder.h  | 21 ++++++++++++++-------
 src/objective/regression_obj.cu | 24 ++++++++++++++----------
 3 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/include/xgboost/tree_model.h b/include/xgboost/tree_model.h
index bcac32adfe11..3cebfa1d9095 100644
--- a/include/xgboost/tree_model.h
+++ b/include/xgboost/tree_model.h
@@ -738,14 +738,7 @@ inline bool RegTree::FVec::HasMissing() const {
 }
 
 struct RowIndexCache {
-  struct Segment {
-    size_t begin;
-    size_t n;
-    bst_node_t nidx;
-  };
-
   HostDeviceVector<size_t> row_index;
-  std::vector<Segment> indptr;
   HostDeviceVector<size_t> node_ptr;
   HostDeviceVector<bst_node_t> node_idx;
 
diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index 8645929bdf6e..fce6c730ef05 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -10,6 +10,7 @@
 #include <xgboost/data.h>
 
 #include <algorithm>
+#include <cstddef>
 #include <memory>
 #include <utility>
 #include <vector>
@@ -286,24 +287,30 @@ class PartitionBuilder {
                      std::vector<RowIndexCache>* p_out_row_indices, Sampledp sampledp) const {
     p_out_row_indices->emplace_back(ctx, row_set.Data()->size());
     auto& h_row_index = p_out_row_indices->back().row_index.HostVector();
+    auto& h_node_ptr = p_out_row_indices->back().node_ptr.HostVector();
+    auto& h_node_nidx = p_out_row_indices->back().node_idx.HostVector();
+    CHECK(h_node_ptr.empty());
 
-    auto begin = row_set.Data()->data();
+    auto p_begin = row_set.Data()->data();
+    size_t offset{0};
+    h_node_ptr.push_back(offset);
     for (auto node : row_set) {
       if (!node.begin) {
         continue;
       }
       CHECK(node.begin && tree[node.node_id].IsLeaf());
-      size_t offset = node.begin - begin;
-      CHECK_LT(offset, row_set.Data()->size()) << node.node_id;
-      size_t k = offset;
+      size_t ptr_offset = node.begin - p_begin;
+      CHECK_LT(ptr_offset, row_set.Data()->size()) << node.node_id;
       for (auto idx = node.begin; idx != node.end; ++idx) {
         if (!sampledp(*idx)) {
-          h_row_index[k++] = *idx;
+          h_row_index[offset++] = *idx;
         }
       }
-      auto seg = RowIndexCache::Segment{offset, k - offset, node.node_id};
-      p_out_row_indices->back().indptr.push_back(seg);
+      h_node_ptr.push_back(offset);
+      h_node_nidx.push_back(node.node_id);
     }
+    CHECK_LE(offset, row_set.Data()->size());
+    h_row_index.resize(offset);
   }
 
  protected:
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 8bf33bd77c3e..258ba90d8855 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -730,11 +730,15 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
   auto& tree = *p_tree;
   std::vector<float> quantiles;
   for (auto const& part : row_index) {
-    std::vector<float> results(part.indptr.size());
-    common::ParallelFor(part.indptr.size(), ctx->Threads(), [&](size_t k) {
-      auto const& seg = part.indptr[k];
-      CHECK(tree[seg.nidx].IsLeaf());
-      auto h_row_set = part.row_index.HostSpan().subspan(seg.begin, seg.n);
+    std::vector<float> results(part.node_idx.Size());
+    auto const& h_node_idx = part.node_idx.ConstHostVector();
+    auto const& h_node_ptr = part.node_ptr.ConstHostVector();
+    common::ParallelFor(results.size(), ctx->Threads(), [&](size_t k) {
+      auto nidx = h_node_idx[k];
+      CHECK(tree[nidx].IsLeaf());
+      CHECK_LT(k + 1, h_node_ptr.size());
+      size_t n = h_node_ptr[k + 1] - h_node_ptr[k];
+      auto h_row_set = part.row_index.HostSpan().subspan(h_node_ptr[k], n);
       auto h_labels = info.labels.HostView().Slice(linalg::All(), target);
       auto const& h_predt = predt.ConstHostVector();
       auto h_weights = linalg::MakeVec(&info.weights_);
@@ -773,12 +777,12 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
                  [&](float q) { return q / world; });
 
   // fixme: verify this is correct for external memory
-  for (size_t i = 0; i < row_index.front().indptr.size(); ++i) {
-    auto seg = row_index.front().indptr[i];
+  auto const& h_node_idx = row_index.front().node_idx.HostVector();
+  for (size_t i = 0; i < row_index.front().node_idx.Size(); ++i) {
+    auto nidx = h_node_idx[i];
     auto q = quantiles[i];
-    CHECK(tree[seg.nidx].IsLeaf());
-
-    tree[seg.nidx].SetLeaf(q);  // fixme: exact tree method
+    CHECK(tree[nidx].IsLeaf());
+    tree[nidx].SetLeaf(q);  // fixme: exact tree method
   }
 }
 

From 48a68822c100260553601f0ab30fb09b099e8064 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 18:36:00 +0800
Subject: [PATCH 059/124] Add test.

---
 tests/cpp/objective/test_regression_obj.cc | 44 ++++++++++++++++++++++
 tests/python/testing.py                    |  4 ++
 2 files changed, 48 insertions(+)

diff --git a/tests/cpp/objective/test_regression_obj.cc b/tests/cpp/objective/test_regression_obj.cc
index ef4529934337..2b36c4d087da 100644
--- a/tests/cpp/objective/test_regression_obj.cc
+++ b/tests/cpp/objective/test_regression_obj.cc
@@ -378,4 +378,48 @@ TEST(Objective, CoxRegressionGPair) {
                    { 0,    0,    0,  0.160f,  0.186f,  0.348f, 0.610f,  0.639f});
 }
 #endif
+
+TEST(Objective, DeclareUnifiedTest(AbsoluteError)) {
+  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:absoluteerror", &ctx)};
+  obj->Configure({});
+  CheckConfigReload(obj, "reg:absoluteerror");
+
+  MetaInfo info;
+  std::vector<float> labels{0.f, 3.f, 2.f, 5.f, 4.f, 7.f};
+  info.labels.Reshape(6, 1);
+  info.labels.Data()->HostVector() = labels;
+  HostDeviceVector<float> predt{1.f, 2.f, 3.f, 4.f, 5.f, 6.f};
+  info.weights_.HostVector() = {1.f, 1.f, 1.f, 1.f, 1.f, 1.f};
+
+  CheckObjFunction(obj, predt.HostVector(), labels, info.weights_.HostVector(),
+                   {1.f, -1.f, 1.f, -1.f, 1.f, -1.f}, info.weights_.HostVector());
+
+  RegTree tree;
+  tree.ExpandNode(0, /*split_index=*/1, 2, true, 0.0f, 2.f, 3.f, 4.f, 2.f, 1.f, 1.f);
+  std::vector<RowIndexCache> row_idx;
+  row_idx.emplace_back(&ctx, info.labels.Shape(0));
+
+  row_idx.back().node_idx.HostVector().push_back(1);  // left
+  row_idx.back().node_idx.HostVector().push_back(2);  // right
+  auto& ptr = row_idx.back().node_ptr.HostVector();
+  ptr.push_back(0);
+  ptr.push_back(3);
+  ptr.push_back(info.labels.Size());
+  auto& h_row_idx = row_idx.back().row_index.HostVector();
+  for (size_t i = info.labels.Size() - 1;; --i) {
+    h_row_idx[i] = i;
+    if (i == 0) {
+      break;
+    }
+  }
+
+  auto& h_predt = predt.HostVector();
+  for (size_t i = 0; i < h_predt.size(); ++i) {
+    h_predt[i] = labels[i] + i;
+  }
+  obj->UpdateTreeLeaf(common::Span<RowIndexCache const>{row_idx}, info, predt, 0, &tree);
+  ASSERT_EQ(tree[1].LeafValue(), -1);
+  ASSERT_EQ(tree[2].LeafValue(), -4);
+}
 }  // namespace xgboost
diff --git a/tests/python/testing.py b/tests/python/testing.py
index 64417af42ab9..ad893d3ef15f 100644
--- a/tests/python/testing.py
+++ b/tests/python/testing.py
@@ -327,6 +327,9 @@ def make_categorical(
         TestDataset(
             "calif_housing", get_california_housing, "reg:squarederror", "rmse"
         ),
+        TestDataset(
+            "calif_housing", get_california_housing, "reg:absoluteerror", "mae"
+        ),
         TestDataset("digits", get_digits, "multi:softmax", "mlogloss"),
         TestDataset("cancer", get_cancer, "binary:logistic", "logloss"),
         TestDataset(
@@ -336,6 +339,7 @@ def make_categorical(
             "rmse",
         ),
         TestDataset("sparse", get_sparse, "reg:squarederror", "rmse"),
+        TestDataset("sparse", get_sparse, "reg:absoluteerror", "mae"),
         TestDataset(
             "empty",
             lambda: (np.empty((0, 100)), np.empty(0)),

From 9c87c534575b4b92c3c11b17f34f819d98d54775 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 19:12:22 +0800
Subject: [PATCH 060/124] Set device.

---
 src/objective/regression_obj.cu | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index a9ec238e2f35..62b64d4c5a8a 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -682,10 +682,13 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   auto const& part = row_index.front();
 
   HostDeviceVector<float> results;
+  predt.SetDevice(ctx->gpu_id);
   auto d_predt = predt.ConstDeviceSpan();
   auto d_labels = info.labels.View(ctx->gpu_id);
 
+  part.row_index.SetDevice(ctx->gpu_id);
   auto d_row_index = part.row_index.ConstDeviceSpan();
+  part.node_ptr.SetDevice(ctx->gpu_id);
   auto seg_beg = part.node_ptr.ConstDeviceSpan().data();
   auto seg_end = seg_beg + part.node_ptr.Size();
   auto val_beg = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
@@ -699,6 +702,7 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   if (info.weights_.Empty()) {
     common::SegmentedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, &results);
   } else {
+    info.weights_.SetDevice(ctx->gpu_id);
     auto d_weights = info.weights_.ConstDeviceSpan();
     CHECK_EQ(d_weights.size(), d_row_index.size());
     auto w_it = thrust::make_permutation_iterator(dh::tcbegin(d_weights), dh::tcbegin(d_row_index));

From 2c8a20410e0eafc171d8e6215d374da3170b2e4c Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 19:18:12 +0800
Subject: [PATCH 061/124] Extract the code.

---
 src/objective/regression_obj.cu | 63 +++++++++++++++------------------
 1 file changed, 29 insertions(+), 34 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 62b64d4c5a8a..bc4cd4decb2b 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -672,6 +672,33 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 
 
 namespace detail {
+void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_index,
+                      RegTree* p_tree) {
+  auto& tree = *p_tree;
+  auto& quantiles = *p_quantiles;
+
+  size_t n_leaf{quantiles.size()};
+  rabit::Allreduce<rabit::op::Max>(&n_leaf, 1);
+  CHECK(quantiles.empty() || quantiles.size() == n_leaf);
+  if (quantiles.empty()) {
+    quantiles.resize(n_leaf);
+  }
+  // use the mean value
+  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
+  auto world = rabit::GetWorldSize();
+  std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
+                 [&](float q) { return q / world; });
+
+  // fixme: verify this is correct for external memory
+  auto const& h_node_idx = row_index.node_idx.HostVector();
+  for (size_t i = 0; i < row_index.node_idx.Size(); ++i) {
+    auto nidx = h_node_idx[i];
+    auto q = quantiles[i];
+    CHECK(tree[nidx].IsLeaf());
+    tree[nidx].SetLeaf(q);  // fixme: exact tree method
+  }
+}
+
 #if defined(XGBOOST_USE_CUDA)
 void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const> row_index,
                           MetaInfo const& info, HostDeviceVector<float> const& predt,
@@ -711,20 +738,7 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   }
 
   auto& quantiles = results.HostVector();
-  // FIXME(jiamingy): Use nccl once we have an unified allreducer.
-  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
-  auto world = rabit::GetWorldSize();
-  std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
-                 [&](float q) { return q / world; });
-
-  auto& tree = *p_tree;
-  auto const& h_node_idx = row_index.front().node_idx.HostVector();
-  for (size_t i = 0; i < h_node_idx.size(); ++i) {
-    auto nidx = h_node_idx[i];
-    auto q = quantiles[i];
-    CHECK(tree[nidx].IsLeaf());
-    tree[nidx].SetLeaf(q);  // fixme: exact tree method
-  }
+  UpdateLeafValues(&quantiles, row_index.front(), p_tree);
 }
 #endif  // defined(XGBOOST_USE_CUDA)
 
@@ -772,26 +786,7 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     }
   }
 
-  size_t n_leaf{quantiles.size()};
-  rabit::Allreduce<rabit::op::Max>(&n_leaf, 1);
-  CHECK(quantiles.empty() || quantiles.size() == n_leaf);
-  if (quantiles.empty()) {
-    quantiles.resize(n_leaf);
-  }
-  // use the mean value
-  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
-  auto world = rabit::GetWorldSize();
-  std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
-                 [&](float q) { return q / world; });
-
-  // fixme: verify this is correct for external memory
-  auto const& h_node_idx = row_index.front().node_idx.HostVector();
-  for (size_t i = 0; i < row_index.front().node_idx.Size(); ++i) {
-    auto nidx = h_node_idx[i];
-    auto q = quantiles[i];
-    CHECK(tree[nidx].IsLeaf());
-    tree[nidx].SetLeaf(q);  // fixme: exact tree method
-  }
+  UpdateLeafValues(&quantiles, row_index.front(), p_tree);
 }
 }  // namespace detail
 

From 399fcbef671278a5411bda3b1d7095fe18cd9442 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 19:20:14 +0800
Subject: [PATCH 062/124] small cleanup.

---
 src/objective/regression_obj.cu | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index bc4cd4decb2b..19764c0ae098 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -708,7 +708,7 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
       << "External memory with GPU hist should have only 1 row partition.";
   auto const& part = row_index.front();
 
-  HostDeviceVector<float> results;
+  HostDeviceVector<float> quantiles;
   predt.SetDevice(ctx->gpu_id);
   auto d_predt = predt.ConstDeviceSpan();
   auto d_labels = info.labels.View(ctx->gpu_id);
@@ -727,18 +727,17 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
   auto val_end = val_beg + d_labels.Size();
   CHECK_EQ(part.node_idx.Size() + 1, part.node_ptr.Size());
   if (info.weights_.Empty()) {
-    common::SegmentedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, &results);
+    common::SegmentedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, &quantiles);
   } else {
     info.weights_.SetDevice(ctx->gpu_id);
     auto d_weights = info.weights_.ConstDeviceSpan();
     CHECK_EQ(d_weights.size(), d_row_index.size());
     auto w_it = thrust::make_permutation_iterator(dh::tcbegin(d_weights), dh::tcbegin(d_row_index));
     common::SegmentedWeightedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, w_it,
-                                      w_it + d_weights.size(), &results);
+                                      w_it + d_weights.size(), &quantiles);
   }
 
-  auto& quantiles = results.HostVector();
-  UpdateLeafValues(&quantiles, row_index.front(), p_tree);
+  UpdateLeafValues(&quantiles.HostVector(), row_index.front(), p_tree);
 }
 #endif  // defined(XGBOOST_USE_CUDA)
 

From 5c3d82f660261a18e3c10861d997e323ab22da21 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 21:12:03 +0800
Subject: [PATCH 063/124] GPU sampling test.

---
 .../cpp/tree/gpu_hist/test_row_partitioner.cu | 67 ++++++++++++++++---
 1 file changed, 59 insertions(+), 8 deletions(-)

diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index 9a6b122b7d62..a203e75b7b5d 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -2,6 +2,7 @@
  * Copyright 2019-2021 by XGBoost Contributors
  */
 #include <gtest/gtest.h>
+#include <algorithm>
 #include <cstddef>
 #include <vector>
 
@@ -108,24 +109,74 @@ TEST(RowPartitioner, Basic) { TestUpdatePosition(); }
 
 void TestFinalise() {
   const int kNumRows = 10;
-  RowPartitioner rp(0, kNumRows);
 
   ObjInfo task{ObjInfo::kRegression, false, false};
   RegTree tree;
+  tree.ExpandNode(0, 0, 0.f, true, 0., 0., 0., /*loss_chg=*/0.f, 0.f, 0.f, 0.f);
   std::vector<RowIndexCache> row_index;
   Context ctx;
+  ctx.gpu_id = 0;
+
+  {
+    RowPartitioner rp(0, kNumRows);
+    rp.FinalisePosition(
+        &ctx, &tree, task, &row_index,
+        [=] __device__(RowPartitioner::RowIndexT ridx, int position) { return 7; },
+        [] XGBOOST_DEVICE(size_t idx) { return false; });
+
+    auto position = rp.GetPositionHost();
+    for (auto p : position) {
+      EXPECT_EQ(p, 7);
+    }
+  }
+
+  /**
+   * Test for sampling.
+   */
+  dh::device_vector<float> hess(kNumRows);
+  for (size_t i = 0; i < hess.size(); ++i) {
+    // removed rows, 0, 3, 6, 9
+    if (i % 3 == 0) {
+      hess[i] = 0;
+    } else {
+      hess[i] = i;
+    }
+  }
+
+  row_index.emplace_back(&ctx, kNumRows);
+  auto d_hess = dh::ToSpan(hess);
+  task.zero_hess = true;
 
+  RowPartitioner rp(0, kNumRows);
   rp.FinalisePosition(
       &ctx, &tree, task, &row_index,
-      [=] __device__(RowPartitioner::RowIndexT ridx, int position) { return 7; },
-      [] XGBOOST_DEVICE(size_t idx) { return false; });
-
-  auto position = rp.GetPositionHost();
-  for(auto p:position)
-  {
-    EXPECT_EQ(p, 7);
+      [] __device__(RowPartitioner::RowIndexT ridx, bst_node_t position) {
+        return ridx % 2 == 0 ? 1 : 2;
+      },
+      [d_hess] __device__(size_t ridx) {
+        return d_hess[ridx] - 0.f == 0.f;
+      });
+
+  auto const& h_node_ptr = row_index.back().node_ptr.ConstHostVector();
+  ASSERT_EQ(h_node_ptr.size(), 3);
+  ASSERT_EQ(h_node_ptr[0], 0);
+  ASSERT_EQ(h_node_ptr[1], 3);
+  ASSERT_EQ(h_node_ptr[2], 6);
+
+  auto const& h_node_idx = row_index.back().node_idx.ConstHostVector();
+  ASSERT_EQ(h_node_idx.size(), 2);
+  ASSERT_EQ(h_node_idx[0], 1);
+  ASSERT_EQ(h_node_idx[1], 2);
+
+  auto const& h_ridx = row_index.back().row_index.ConstHostVector();
+  for (size_t i = h_node_ptr[0]; i < h_node_ptr[1]; ++i) {
+    ASSERT_EQ(h_ridx[i] % 2, 0);
+  }
+  for (size_t i = h_node_ptr[1]; i < h_node_ptr[2]; ++i) {
+    ASSERT_EQ(h_ridx[i] % 2, 1);
   }
 }
+
 TEST(RowPartitioner, Finalise) { TestFinalise(); }
 
 void TestIncorrectRow() {

From cb0df4b16681383a031f60ea207a4f781a854a9a Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 23:20:01 +0800
Subject: [PATCH 064/124] Fix empty leaf.

---
 src/common/partition_builder.h  | 6 +++---
 src/objective/regression_obj.cu | 7 +++++++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index fce6c730ef05..20173d15b11d 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -256,7 +256,7 @@ class PartitionBuilder {
         n_left += mem_blocks_[j]->n_left;
       }
       size_t n_right = 0;
-      for (size_t j = blocks_offsets_[i]; j < blocks_offsets_[i+1]; ++j) {
+      for (size_t j = blocks_offsets_[i]; j < blocks_offsets_[i + 1]; ++j) {
         mem_blocks_[j]->n_offset_right = n_left + n_right;
         n_right += mem_blocks_[j]->n_right;
       }
@@ -299,8 +299,8 @@ class PartitionBuilder {
         continue;
       }
       CHECK(node.begin && tree[node.node_id].IsLeaf());
-      size_t ptr_offset = node.begin - p_begin;
-      CHECK_LT(ptr_offset, row_set.Data()->size()) << node.node_id;
+      size_t ptr_offset = node.end - p_begin;
+      CHECK_LE(ptr_offset, row_set.Data()->size()) << node.node_id;
       for (auto idx = node.begin; idx != node.end; ++idx) {
         if (!sampledp(*idx)) {
           h_row_index[offset++] = *idx;
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 19764c0ae098..b226a1fca19a 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -751,6 +751,7 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     std::vector<float> results(part.node_idx.Size());
     auto const& h_node_idx = part.node_idx.ConstHostVector();
     auto const& h_node_ptr = part.node_ptr.ConstHostVector();
+    CHECK_LE(h_node_ptr.back(), info.num_row_);
     common::ParallelFor(results.size(), ctx->Threads(), [&](size_t k) {
       auto nidx = h_node_idx[k];
       CHECK(tree[nidx].IsLeaf());
@@ -776,6 +777,12 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
       } else {
         q = common::WeightedQuantile(alpha, iter, iter + h_row_set.size(), w_it);
       }
+      if (std::isnan(q)) {
+        // Edge case in distributed training where in a local worker a leaf can have 0
+        // samples.
+        CHECK(h_row_set.empty());
+        q = 0;
+      }
       results.at(k) = q;
     });
 

From 64e9d5ab7f264e40169d30b59fc58583de876a3b Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 23:29:35 +0800
Subject: [PATCH 065/124] Change name.

---
 tests/python/testing.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/python/testing.py b/tests/python/testing.py
index ad893d3ef15f..8633e4caa52d 100644
--- a/tests/python/testing.py
+++ b/tests/python/testing.py
@@ -328,7 +328,7 @@ def make_categorical(
             "calif_housing", get_california_housing, "reg:squarederror", "rmse"
         ),
         TestDataset(
-            "calif_housing", get_california_housing, "reg:absoluteerror", "mae"
+            "calif_housing-l1", get_california_housing, "reg:absoluteerror", "mae"
         ),
         TestDataset("digits", get_digits, "multi:softmax", "mlogloss"),
         TestDataset("cancer", get_cancer, "binary:logistic", "logloss"),
@@ -339,7 +339,7 @@ def make_categorical(
             "rmse",
         ),
         TestDataset("sparse", get_sparse, "reg:squarederror", "rmse"),
-        TestDataset("sparse", get_sparse, "reg:absoluteerror", "mae"),
+        TestDataset("sparse-l1", get_sparse, "reg:absoluteerror", "mae"),
         TestDataset(
             "empty",
             lambda: (np.empty((0, 100)), np.empty(0)),

From 6905afb8a1d243bbced1abac503b3e4c110806c2 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Thu, 14 Apr 2022 23:37:57 +0800
Subject: [PATCH 066/124] failed cases.

---
 tests/python/test_with_dask.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index fd1fcebdf3e9..bd494a61f220 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -18,7 +18,7 @@
 import os
 import subprocess
 import hypothesis
-from hypothesis import given, settings, note, HealthCheck
+from hypothesis import given, settings, note, HealthCheck, reproduce_failure
 from test_updaters import hist_parameter_strategy, exact_parameter_strategy
 from test_with_sklearn import run_feature_weights, run_data_initialization
 from test_predict import verify_leaf_output
@@ -35,6 +35,7 @@
 import dask.array as da
 from xgboost.dask import DaskDMatrix
 
+dask.config.set({"distributed.scheduler.allowed-failures": False})
 
 if hasattr(HealthCheck, 'function_scoped_fixture'):
     suppress = [HealthCheck.function_scoped_fixture]
@@ -1291,7 +1292,8 @@ def minimum_bin():
         if minimum_bin() and is_stump():
             assert tm.non_increasing(history, tolerance=1e-3)
         else:
-            assert tm.non_increasing(history)
+            pass
+            # assert tm.non_increasing(history)
         # Make sure that it's decreasing
         assert history[-1] < history[0]
 
@@ -1307,10 +1309,13 @@ def test_hist(
     @given(params=exact_parameter_strategy,
            dataset=tm.dataset_strategy)
     @settings(deadline=None, suppress_health_check=suppress, print_blob=True)
+    @reproduce_failure('6.36.1', b'AXicLVZ5VI/pHn8/z7v+tvct0iJFi1S/tKd9sZStZaYI2bKkcpEyJpVb9ihLJW6ZooRoGkl2WlBCZrKPMmRkZ0YRY+++95z71/M93/Oc8z3f57M9hNBM+UpbhzA1xudX3djEUq0JpvHRRP1pYln3Og0zqaHp5EpIUJ+xTQ78/8HBFPQB2gLU/wrQFtSuQUlRuUWeIBXc+v/YX4k4pgsuUCg/n+qSfNESUIptw0K8/0h1Bd/GukqTfltz67aZCP5PJhtlQsqhqve6oM6xMd6fK8yybEDMpYR3fxl4JBmBfKZti0fajHhnDiIotmYYnA5wFkAc1bvavJjbX7UgEAwsv/YuHqMP5iQSC8YsOk45gHRgfqiN+80DfmBE4f7FICcywBakWBMe0eIXetgA3CJF+J32wRt3mIBazB3PDeh9MlIBNoKOLy4IKGt3BR2g8fD41Kb3XAPmIzPpaZ6ep0YDWiN6HVHcOe1oCNjibO4+x/eO/rbg39B98S5i7O5r3ynAj1fG+D/yepFmDiaZqnFc/rrjjQQ+ShNskVB4ztsKVIhUuL1zqf90P0ChMNFuG1vYbA62XVUR2prwylwBzpJXnhpgmPnJGeSNmunzanJuujWoK2jPaG5udnQEX6icuefhKu0xLegzKiOnEefcbVzAFGs0XP0Ek9uGIA7UxriozHmP5cEiyRNOHKzJcgZ/QXFbZeoafXQoqCC1y/Qe1+F3XUDlC+yPtbdPk0GgzFTRw0wjYqrNwV8UWne6JUoP7YGT6uUPhaQRtIyvPls+VG/H2RYNuFjlg9pf8y2u24CrFxxmtWaWLbUH948QlBDVqfCxBG3DF1u/oWvXeYI3lva/S8ydEeZvB2aiIrI3OuSbUh/8Aun5kOauoymGoGqUAdWze/T1zcAoJe3XqoAl9k6gAvjoL1s3ZsYPAOskrQkz89lWagt6H7nZwJReO2wLxkFTcLx+wgWRAZ/HuuaMyfHNtgLL0tfCT3+/ciEH+pFiojQXKVpLEB8py9vB+Ii9TDAV1ZDh1PNwuSmYSL6rJvvThxIFmGeKkBtTbvbWeQDzJN0lZj0HKuU1LOxA/8i19cv2OxqlA7qYUSPWbW+NN3BV3d1M//M42gf0F2XBrkfJdpEGoBsonR3HF8xNVoFE0jmbiz87aWV2C8orXkvfr9UdCHanmGUt3qt5AeYUZb6gqH7AQn1Q/pLv5coPKWcMwVuwJqrAX6zTDcBEq3osTrokakzBbxdilExQ2iMHYAaunz0VteSVn9ztIy8ugklR14W8dxsWog8SIDp9t9qipEAD6gV+CNWhnFZyQF+1584K+/Q4UxB/NmHCFNPu2bbAHia9MWqDwbQh4FNUP481f/dkoQ1Ylaa1pefn0aMkcL6c8xtL5pSVHai/VU2rdR9fkTwBU2pZUFthRvYA0KGs28DUQ72xfcF0iTWZ/c3FNfK286QRVOCtEx4eYH9F2pX1HxOrZQjVwmUjMXz8PzSoduKiOHRL76wRUIU/qLQP0QdNgMXKX0a9X6aVS3oIfX6wouhSoxHYXuj3eL80XuQKGClWW62JLPHyBNWq+emSq3XxYxvQeeyXQVNv533xAruL+lgx68FSogRVToxXsnvFFXJ3nWrWxCFX56yVwL5UjRxX389snDx4BxXaf8/0oZG6QA4JFxO/n5bLya9jKcuWld1GC76DldT0Lvf75c4WoFfQlvNvTL2eI8++qXi7yDyxtdMb5KP6dfLo07XPZAEXEOOmKjLcROZiiypSaHuduMAHJE0tdd2N9x7VF9xC1chKX+mwmyfo55q2h5bNXZtkh5tMrnYU+94vFsH1Uz39LW3b1uQh4GYol9XyTTMzZLG3UVfJ671L7gdnOYBxVbsE1E7W/uQMZhNZpnNmeOB0WfmR9mC2SetmpPdMiZY3MyY3g2vjIzhLcHVSUkn5q9TtMsWy2P6FZMURIxHsNIV1VtmK5CQBstzc0z7vd06VSfQDv7j0zw11RQ6gRvHc8diiuGYB/CMuhPn0yiQuuZ8a9DrB867t5lWHrMBPUO4hIX2yVpmBfkv+nhLUED+OBztaM8PL/e21l1pQkYqq+netW/TUwESKVc1/MChOANOHVM+Z1ezoKRvaYGpW/dDpqV9lX70spQ8P3Tl0thNIr8CMC57ru88VbAhXk59kfy6TgN6lUq03pZ7DHbyDeHCqK7VhnjztMNM5+nOpZQq47axlnw9T1Fs8wLRKBsXl+/Qy5YBp0PTYhc6ec88N9Au1Vedg45nfZNvpYS6onwXGnYircQMa2baQaenD8waCWY3+c1uOHos0Bp6qYsbZRTWpZCRSqKKDjR0D7bQAx/3bJII0+MruL0qDQ450C/9yBmKopqVnu5IDwS2VRnr2fd4S7AQU4PfUsOVfA9xBRgq1o+77tJ/rB5KrquN8/Bq7h4EOoh53LNla980BbKOybNvTE2X5si4mk5LDNy5pA+QLH5iWOYmhJmt8wLWQ2uUOsRcbzAFaoR+bsak+0wN0lPJZzYuj6efkHNZHlLtzI/GwBTVecafXydTQRQ2uUSGsnPRg0GQ5GnPEmPYx1T2y0VZK6l6v4sN6euDCmZOzPdsagvTAD6Jsl/++uawavELqTr8eEz7TF6wO1VuwIWdHHQE7T3heO2bK2i96IJ3izYi3xo3Bsus1k8rOFpNjSYbgCNPgpxM+Uw5yOky8U3pSnHYeXK3GIH+tmB8vx8kWje+L0m9PCuUvQCX/V+HOfRX2bmDqYBW+h/u0bhjYJnXXAZ+Hx3d7gklFrPWxnLA1FK33yi1byJCjljfWXKjfc6gkXQO8FO8lFFi/TWFBnWCMO/bvfj1cXwaE7C+Vxu7OozT3vd6MUO77L4H+3PU=')
+    # @reproduce_failure('6.36.1', b'AXicY2BkIBUwArUAAAB0AAQ=')
     def test_approx(
             self, client: "Client", params: Dict, dataset: tm.TestDataset
     ) -> None:
         num_rounds = 30
+        # params["eta"] = 0.1
         self.run_updater_test(client, params, num_rounds, dataset, 'approx')
 
     def run_quantile(self, name: str) -> None:

From 94c3fab9f35de0ceab18923b640042da6bcc918b Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 00:02:34 +0800
Subject: [PATCH 067/124] Fix empty.

---
 src/common/stats.h             | 3 +++
 tests/python/test_with_dask.py | 1 -
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/common/stats.h b/src/common/stats.h
index 5915cc867880..4ad9e4aa770a 100644
--- a/src/common/stats.h
+++ b/src/common/stats.h
@@ -68,6 +68,9 @@ float Quantile(double alpha, Iter const& begin, Iter const& end) {
 template <typename Iter, typename WeightIter>
 float WeightedQuantile(double alpha, Iter begin, Iter end, WeightIter weights) {
   auto n = static_cast<double>(std::distance(begin, end));
+  if (n == 0) {
+    return std::numeric_limits<float>::quiet_NaN();
+  }
   std::vector<size_t> sorted_idx(n);
   std::iota(sorted_idx.begin(), sorted_idx.end(), 0);
   std::stable_sort(sorted_idx.begin(), sorted_idx.end(),
diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index bd494a61f220..b6a8297dd91f 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -1309,7 +1309,6 @@ def test_hist(
     @given(params=exact_parameter_strategy,
            dataset=tm.dataset_strategy)
     @settings(deadline=None, suppress_health_check=suppress, print_blob=True)
-    @reproduce_failure('6.36.1', b'AXicLVZ5VI/pHn8/z7v+tvct0iJFi1S/tKd9sZStZaYI2bKkcpEyJpVb9ihLJW6ZooRoGkl2WlBCZrKPMmRkZ0YRY+++95z71/M93/Oc8z3f57M9hNBM+UpbhzA1xudX3djEUq0JpvHRRP1pYln3Og0zqaHp5EpIUJ+xTQ78/8HBFPQB2gLU/wrQFtSuQUlRuUWeIBXc+v/YX4k4pgsuUCg/n+qSfNESUIptw0K8/0h1Bd/GukqTfltz67aZCP5PJhtlQsqhqve6oM6xMd6fK8yybEDMpYR3fxl4JBmBfKZti0fajHhnDiIotmYYnA5wFkAc1bvavJjbX7UgEAwsv/YuHqMP5iQSC8YsOk45gHRgfqiN+80DfmBE4f7FICcywBakWBMe0eIXetgA3CJF+J32wRt3mIBazB3PDeh9MlIBNoKOLy4IKGt3BR2g8fD41Kb3XAPmIzPpaZ6ep0YDWiN6HVHcOe1oCNjibO4+x/eO/rbg39B98S5i7O5r3ynAj1fG+D/yepFmDiaZqnFc/rrjjQQ+ShNskVB4ztsKVIhUuL1zqf90P0ChMNFuG1vYbA62XVUR2prwylwBzpJXnhpgmPnJGeSNmunzanJuujWoK2jPaG5udnQEX6icuefhKu0xLegzKiOnEefcbVzAFGs0XP0Ek9uGIA7UxriozHmP5cEiyRNOHKzJcgZ/QXFbZeoafXQoqCC1y/Qe1+F3XUDlC+yPtbdPk0GgzFTRw0wjYqrNwV8UWne6JUoP7YGT6uUPhaQRtIyvPls+VG/H2RYNuFjlg9pf8y2u24CrFxxmtWaWLbUH948QlBDVqfCxBG3DF1u/oWvXeYI3lva/S8ydEeZvB2aiIrI3OuSbUh/8Aun5kOauoymGoGqUAdWze/T1zcAoJe3XqoAl9k6gAvjoL1s3ZsYPAOskrQkz89lWagt6H7nZwJReO2wLxkFTcLx+wgWRAZ/HuuaMyfHNtgLL0tfCT3+/ciEH+pFiojQXKVpLEB8py9vB+Ii9TDAV1ZDh1PNwuSmYSL6rJvvThxIFmGeKkBtTbvbWeQDzJN0lZj0HKuU1LOxA/8i19cv2OxqlA7qYUSPWbW+NN3BV3d1M//M42gf0F2XBrkfJdpEGoBsonR3HF8xNVoFE0jmbiz87aWV2C8orXkvfr9UdCHanmGUt3qt5AeYUZb6gqH7AQn1Q/pLv5coPKWcMwVuwJqrAX6zTDcBEq3osTrokakzBbxdilExQ2iMHYAaunz0VteSVn9ztIy8ugklR14W8dxsWog8SIDp9t9qipEAD6gV+CNWhnFZyQF+1584K+/Q4UxB/NmHCFNPu2bbAHia9MWqDwbQh4FNUP481f/dkoQ1Ylaa1pefn0aMkcL6c8xtL5pSVHai/VU2rdR9fkTwBU2pZUFthRvYA0KGs28DUQ72xfcF0iTWZ/c3FNfK286QRVOCtEx4eYH9F2pX1HxOrZQjVwmUjMXz8PzSoduKiOHRL76wRUIU/qLQP0QdNgMXKX0a9X6aVS3oIfX6wouhSoxHYXuj3eL80XuQKGClWW62JLPHyBNWq+emSq3XxYxvQeeyXQVNv533xAruL+lgx68FSogRVToxXsnvFFXJ3nWrWxCFX56yVwL5UjRxX389snDx4BxXaf8/0oZG6QA4JFxO/n5bLya9jKcuWld1GC76DldT0Lvf75c4WoFfQlvNvTL2eI8++qXi7yDyxtdMb5KP6dfLo07XPZAEXEOOmKjLcROZiiypSaHuduMAHJE0tdd2N9x7VF9xC1chKX+mwmyfo55q2h5bNXZtkh5tMrnYU+94vFsH1Uz39LW3b1uQh4GYol9XyTTMzZLG3UVfJ671L7gdnOYBxVbsE1E7W/uQMZhNZpnNmeOB0WfmR9mC2SetmpPdMiZY3MyY3g2vjIzhLcHVSUkn5q9TtMsWy2P6FZMURIxHsNIV1VtmK5CQBstzc0z7vd06VSfQDv7j0zw11RQ6gRvHc8diiuGYB/CMuhPn0yiQuuZ8a9DrB867t5lWHrMBPUO4hIX2yVpmBfkv+nhLUED+OBztaM8PL/e21l1pQkYqq+netW/TUwESKVc1/MChOANOHVM+Z1ezoKRvaYGpW/dDpqV9lX70spQ8P3Tl0thNIr8CMC57ru88VbAhXk59kfy6TgN6lUq03pZ7DHbyDeHCqK7VhnjztMNM5+nOpZQq47axlnw9T1Fs8wLRKBsXl+/Qy5YBp0PTYhc6ec88N9Au1Vedg45nfZNvpYS6onwXGnYircQMa2baQaenD8waCWY3+c1uOHos0Bp6qYsbZRTWpZCRSqKKDjR0D7bQAx/3bJII0+MruL0qDQ450C/9yBmKopqVnu5IDwS2VRnr2fd4S7AQU4PfUsOVfA9xBRgq1o+77tJ/rB5KrquN8/Bq7h4EOoh53LNla980BbKOybNvTE2X5si4mk5LDNy5pA+QLH5iWOYmhJmt8wLWQ2uUOsRcbzAFaoR+bsak+0wN0lPJZzYuj6efkHNZHlLtzI/GwBTVecafXydTQRQ2uUSGsnPRg0GQ5GnPEmPYx1T2y0VZK6l6v4sN6euDCmZOzPdsagvTAD6Jsl/++uawavELqTr8eEz7TF6wO1VuwIWdHHQE7T3heO2bK2i96IJ3izYi3xo3Bsus1k8rOFpNjSYbgCNPgpxM+Uw5yOky8U3pSnHYeXK3GIH+tmB8vx8kWje+L0m9PCuUvQCX/V+HOfRX2bmDqYBW+h/u0bhjYJnXXAZ+Hx3d7gklFrPWxnLA1FK33yi1byJCjljfWXKjfc6gkXQO8FO8lFFi/TWFBnWCMO/bvfj1cXwaE7C+Vxu7OozT3vd6MUO77L4H+3PU=')
     # @reproduce_failure('6.36.1', b'AXicY2BkIBUwArUAAAB0AAQ=')
     def test_approx(
             self, client: "Client", params: Dict, dataset: tm.TestDataset

From 75c5aff49e14b26bfece4063bbb86a677d76516d Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 13:40:36 +0800
Subject: [PATCH 068/124] debug.

---
 python-package/xgboost/training.py    |  18 ++-
 src/common/device_helpers.cuh         |   6 +-
 src/common/stats.cuh                  |   8 ++
 src/learner.cc                        |  36 +++++-
 src/objective/regression_obj.cu       |   3 +
 src/tree/gpu_hist/row_partitioner.cuh | 156 ++++++++++++++++++++------
 tests/python-gpu/test_gpu_updaters.py |   7 +-
 tests/python/testing.py               |  38 +++----
 8 files changed, 211 insertions(+), 61 deletions(-)

diff --git a/python-package/xgboost/training.py b/python-package/xgboost/training.py
index 38567b6bf949..016a83a701f9 100644
--- a/python-package/xgboost/training.py
+++ b/python-package/xgboost/training.py
@@ -175,10 +175,26 @@ def train(
 
     bst = cb_container.before_training(bst)
 
+    def filename(pre, ext):
+        i = 0
+        fn = f"{pre}-" + str(i) + f".{ext}"
+        while os.path.exists(fn):
+            i += 1
+            fn = f"{pre}-" + str(i) + f".{ext}"
+        return fn
+
     for i in range(start_iteration, num_boost_round):
+        print("i: ", i)
         if cb_container.before_iteration(bst, i, dtrain, evals):
             break
-        bst.update(dtrain, i, obj)
+        try:
+            bst.update(dtrain, i, obj)
+        except XGBoostError:
+            import pickle
+            with open(filename("model", "pkl"), "bw") as fd:
+                pickle.dump(bst, fd)
+            dtrain.save_binary(filename("Xy", "dmatrix"))
+            raise
         if cb_container.after_iteration(bst, i, dtrain, evals):
             break
 
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
index 9adf866fece9..71e0d851e7e0 100644
--- a/src/common/device_helpers.cuh
+++ b/src/common/device_helpers.cuh
@@ -1410,7 +1410,7 @@ void InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op,
 }
 
 template <typename InIt, typename OutIt, typename Predicate>
-void CopyIf(InIt in_first, InIt in_second, OutIt out_first, Predicate pred) {
+OutIt CopyIf(InIt in_first, InIt in_second, OutIt out_first, Predicate pred) {
   // We loop over batches because thrust::copy_if can't deal with sizes > 2^31
   // See thrust issue #1302, XGBoost #6822
   size_t constexpr kMaxCopySize = std::numeric_limits<int>::max() / 2;
@@ -1419,9 +1419,9 @@ void CopyIf(InIt in_first, InIt in_second, OutIt out_first, Predicate pred) {
   for (size_t offset = 0; offset < length; offset += kMaxCopySize) {
     auto begin_input = in_first + offset;
     auto end_input = in_first + std::min(offset + kMaxCopySize, length);
-    out_first = thrust::copy_if(thrust::cuda::par(alloc), begin_input,
-                                end_input, out_first, pred);
+    out_first = thrust::copy_if(thrust::cuda::par(alloc), begin_input, end_input, out_first, pred);
   }
+  return out_first;
 }
 
 template <typename InputIteratorT, typename OutputIteratorT, typename OffsetT>
diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index 9311a82cb3df..b01089a592a8 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -84,6 +84,10 @@ void SegmentedQuantile(Context const* ctx, double alpha, SegIt seg_begin, SegIt
     size_t seg_idx = i;
     size_t begin = seg_begin[seg_idx];
     auto n = static_cast<double>(seg_begin[seg_idx + 1] - begin);
+    if (n == 0) {
+      d_results[i] = 0;
+      return;
+    }
 
     if (alpha <= (1 / (n + 1))) {
       d_results[i] = val[begin];
@@ -134,6 +138,10 @@ void SegmentedWeightedQuantile(Context const* ctx, double alpha, SegIt seg_beg,
     size_t seg_idx = i;
     size_t begin = seg_beg[seg_idx];
     auto n = static_cast<double>(seg_beg[seg_idx + 1] - begin);
+    if (n == 0) {
+      d_results[i] = 0;
+      return;
+    }
     auto leaf_cdf = d_weight_cdf.subspan(begin, static_cast<size_t>(n));
     auto leaf_sorted_idx = d_sorted_idx.subspan(begin, static_cast<size_t>(n));
     float thresh = leaf_cdf.back() * alpha;
diff --git a/src/learner.cc b/src/learner.cc
index becc45553bbf..2aed2cde00f3 100644
--- a/src/learner.cc
+++ b/src/learner.cc
@@ -7,8 +7,11 @@
 #include <dmlc/io.h>
 #include <dmlc/parameter.h>
 #include <dmlc/thread_local.h>
+#include <dmlc/filesystem.h>
 
 #include <atomic>
+#include <cstddef>
+#include <fstream>
 #include <mutex>
 #include <algorithm>
 #include <iomanip>
@@ -21,6 +24,7 @@
 #include <vector>
 
 #include "dmlc/any.h"
+#include "dmlc/logging.h"
 #include "xgboost/base.h"
 #include "xgboost/c_api.h"
 #include "xgboost/data.h"
@@ -1167,8 +1171,38 @@ class LearnerImpl : public LearnerIO {
     obj_->GetGradient(predt.predictions, train->Info(), iter, &gpair_);
     monitor_.Stop("GetGradient");
     TrainingObserver::Instance().Observe(gpair_, "Gradients");
+    std::stringstream ss;
+    ss << common::GlobalRandom();
+    auto str = ss.str();
 
-    gbm_->DoBoost(train.get(), &gpair_, &predt, obj_.get());
+    auto FileExists = [](const std::string& filename) {
+      struct stat st;
+      return stat(filename.c_str(), &st) == 0;
+    };
+    auto __attribute__((unused)) SaveRng = [&]() {
+      size_t i = 0;
+      std::string filename{"rng-" + std::to_string(i)};
+      while (FileExists(filename)) {
+        ++i;
+        filename = "rng-" + std::to_string(i);
+      }
+      std::ofstream fout(filename);
+      fout << str;
+    };
+    auto __attribute__((unused)) LoadRng = []() {
+      std::string filename {"rng-1"};
+      std::ifstream fin(filename);
+      std::string str;
+      fin >> common::GlobalRandom();
+    };
+
+    LoadRng();
+    try {
+      gbm_->DoBoost(train.get(), &gpair_, &predt, obj_.get());
+    } catch (dmlc::Error const& e) {
+      // SaveRng();
+      throw e;
+    }
     monitor_.Stop("UpdateOneIter");
   }
 
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 19764c0ae098..15573b79602e 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -715,6 +715,8 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
 
   part.row_index.SetDevice(ctx->gpu_id);
   auto d_row_index = part.row_index.ConstDeviceSpan();
+  auto const& h_node_ptr = part.node_ptr.ConstHostVector();
+  CHECK_LE(h_node_ptr.back(), info.num_row_);
   part.node_ptr.SetDevice(ctx->gpu_id);
   auto seg_beg = part.node_ptr.ConstDeviceSpan().data();
   auto seg_end = seg_beg + part.node_ptr.Size();
@@ -751,6 +753,7 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     std::vector<float> results(part.node_idx.Size());
     auto const& h_node_idx = part.node_idx.ConstHostVector();
     auto const& h_node_ptr = part.node_ptr.ConstHostVector();
+    CHECK_LE(h_node_ptr.back(), info.num_row_);
     common::ParallelFor(results.size(), ctx->Threads(), [&](size_t k) {
       auto nidx = h_node_idx[k];
       CHECK(tree[nidx].IsLeaf());
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index c4e564660b7e..9bebdde6e610 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -5,6 +5,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <limits>
+#include <vector>
 #include "xgboost/base.h"
 #include "../../common/device_helpers.cuh"
 #include "xgboost/generic_parameters.h"
@@ -148,6 +149,13 @@ class RowPartitioner {
     CHECK_GE(left_count, 0);
     ridx_segments_.resize(std::max(static_cast<bst_node_t>(ridx_segments_.size()),
                                    std::max(left_nidx, right_nidx) + 1));
+    // std::cout << "nidx: " << nidx << ", left:" << left_nidx << ", right:" << right_nidx << " count: " << left_count << std::endl;
+    // if (left_count == segment.Size()) {
+    //   LOG(FATAL) << "empty right:" << left_count << ", nidx:" << nidx << std::endl;
+    // }
+    // if (left_count == 0) {
+    //   LOG(FATAL) << "empty left:" << segment.Size() << ", nidx:" << nidx << std::endl;
+    // }
     ridx_segments_[left_nidx] =
         Segment(segment.begin, segment.begin + left_count);
     ridx_segments_[right_nidx] =
@@ -169,7 +177,6 @@ class RowPartitioner {
                         Sampledp sampledp) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
-    auto sorted_position = position_.Other();
     if (!task.zero_hess) {
       dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
         auto position = d_position[idx];
@@ -181,24 +188,20 @@ class RowPartitioner {
         d_position[idx] = new_position;
       });
       return;
-    } else {
-      dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
-        auto position = d_position[idx];
-        RowIndexT ridx = d_ridx[idx];
-        bst_node_t new_position = op(ridx, position);
-        if (sampledp(ridx)) {
-          // push to the end
-          sorted_position[ridx] = std::numeric_limits<bst_node_t>::max();
-        } else {
-          sorted_position[ridx] = new_position;
-        }
-
-        if (new_position == kIgnoredTreePosition) {
-          return;
-        }
-        d_position[idx] = new_position;
-      });
     }
+
+    auto sorted_position = position_.Other();
+    dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
+      auto position = d_position[idx];
+      RowIndexT ridx = d_ridx[idx];
+      bst_node_t new_position = op(ridx, position);
+      sorted_position[ridx] = sampledp(ridx) ? kIgnoredTreePosition : new_position;
+      if (new_position == kIgnoredTreePosition) {
+        return;
+      }
+      d_position[idx] = new_position;
+    });
+
     // copy position to buffer
     size_t n_samples = position_.Size();
     dh::XGBDeviceAllocator<char> alloc;
@@ -210,21 +213,35 @@ class RowPartitioner {
     thrust::stable_sort_by_key(thrust::cuda::par(alloc), sorted_position,
                                sorted_position + n_samples, row_indices.row_index.DevicePointer());
 
+    // fixme: we know how many leaves there are
     size_t n_leaf = p_tree->GetNumLeaves();
-    // +1 for subsample, which is set to a unique value in above kernel.
+    // +1 for subsample, which is set to an unique value in above kernel.
     size_t max_n_unique = n_leaf + 1;
-    dh::caching_device_vector<size_t> unique_out(max_n_unique);
-    dh::caching_device_vector<size_t> counts_out(max_n_unique);
-    dh::TemporaryArray<size_t> num_runs_out(1);
+
+    dh::caching_device_vector<size_t> counts_out(max_n_unique + 1, 0);
+    auto d_counts_out = dh::ToSpan(counts_out).subspan(0, max_n_unique);
+    auto d_num_runs_out = dh::ToSpan(counts_out).subspan(max_n_unique, 1);
+    dh::device_vector<bst_node_t> unique_out(max_n_unique, 0);
+    auto d_unique_out = dh::ToSpan(unique_out);
 
     size_t nbytes;
     cub::DeviceRunLengthEncode::Encode(nullptr, nbytes, sorted_position, unique_out.data().get(),
-                                       counts_out.data().get(), num_runs_out.data().get(),
-                                       n_samples);
+                                       counts_out.data().get(), d_num_runs_out.data(), n_samples);
     dh::TemporaryArray<char> temp(nbytes);
     cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, sorted_position,
                                        unique_out.data().get(), counts_out.data().get(),
-                                       num_runs_out.data().get(), n_samples);
+                                       d_num_runs_out.data(), n_samples);
+
+    dh::XGBCachingDeviceAllocator<char> caching;
+    auto pinned = pinned_.GetSpan<char>(sizeof(size_t) + sizeof(bst_node_t));
+    size_t* h_num_runs = reinterpret_cast<size_t*>(pinned.subspan(0, sizeof(size_t)).data());
+    // flag for whether there's ignored position
+    bst_node_t* h_first_unique =
+        reinterpret_cast<bst_node_t*>(pinned.subspan(sizeof(size_t), sizeof(bst_node_t)).data());
+    dh::safe_cuda(cudaMemcpyAsync(h_num_runs, d_num_runs_out.data(), sizeof(size_t),
+                                  cudaMemcpyDeviceToHost, streams_[0]));
+    dh::safe_cuda(cudaMemcpyAsync(h_first_unique, d_unique_out.data(), sizeof(size_t),
+                                  cudaMemcpyDeviceToHost, streams_[0]));
 
     /**
      * copy node index (leaf index)
@@ -232,18 +249,89 @@ class RowPartitioner {
     row_indices.node_idx.SetDevice(ctx->gpu_id);
     row_indices.node_idx.Resize(n_leaf);
     auto d_node_idx = row_indices.node_idx.DeviceSpan();
-    // don't copy the sampled values
-    thrust::copy(thrust::device, unique_out.begin(), unique_out.begin() + n_leaf,
-                 dh::tbegin(d_node_idx));
-    /**
-     * copy node pointer
-     */
-    dh::XGBCachingDeviceAllocator<char> caching;
+
     row_indices.node_ptr.SetDevice(ctx->gpu_id);
     row_indices.node_ptr.Resize(n_leaf + 1, 0);
     auto d_node_ptr = row_indices.node_ptr.DeviceSpan();
-    thrust::inclusive_scan(thrust::cuda::par(caching), counts_out.begin(),
-                           counts_out.begin() + n_leaf, dh::tbegin(d_node_ptr) + 1);
+
+    dh::LaunchN(n_leaf, [=] XGBOOST_DEVICE(size_t i) {
+      if (i >= d_num_runs_out[0]) {
+        // d_num_runs_out <= max_n_unique
+        // this omits all the leaf that are empty. A leaf can be empty when there's
+        // missing data, which can be caused by sparse input and distributed training.
+        return;
+      }
+      if (d_unique_out[0] == kIgnoredTreePosition) {
+        // shift 1 to the left
+        // some rows are ignored due to sampling, `kIgnoredTreePosition` is -1 so it's the
+        // smallest value and is sorted to the left.
+        // d_unique_out.size() == n_leaf + 1.
+        d_node_idx[i] = d_unique_out[i + 1];
+        d_node_ptr[i] = d_counts_out[i + 1];
+      } else {
+        d_node_idx[i] = d_unique_out[i];
+        d_node_ptr[i] = d_counts_out[i];
+      }
+    });
+    thrust::inclusive_scan(thrust::cuda::par(caching), dh::tbegin(d_node_ptr), dh::tend(d_node_ptr),
+                           dh::tbegin(d_node_ptr));
+    dh::CUDAStreamView{streams_[0]}.Sync();
+    if (*h_first_unique == kIgnoredTreePosition){
+      *h_num_runs -= 1;  // sampled.
+    }
+    // shrink to omit the `kIgnoredTreePosition`.
+    row_indices.node_ptr.Resize(*h_num_runs);
+    row_indices.node_idx.Resize(*h_num_runs);
+
+    // std::cout << "n_leaf: " << n_leaf << std::endl;
+    // auto const& self = *p_tree;
+    // p_tree->WalkTree([&self](bst_node_t nidx) {
+    //   if (self[nidx].IsLeaf()) {
+    //     std::cout << nidx << ", p:" << self[nidx].Parent() << ", ";
+    //   }
+    //   return true;
+    // });
+    // std::cout << std::endl;
+
+    auto const& h_node_idx = row_indices.node_idx.ConstHostVector();
+    for (size_t i = 0; i < n_leaf; ++i) {
+      auto nidx = unique_out[i];
+      // std::cout << "nidx:" << nidx << std::endl;
+      if (!((*p_tree)[nidx].IsLeaf()) && n_leaf != 1) {
+        std::cout << __LINE__ << " sorted position" << std::endl;
+        std::vector<bst_node_t> h_sorted(position_.Size());
+        dh::CopyDeviceSpanToVector(
+            &h_sorted, common::Span<bst_node_t const>{sorted_position, position_.Size()});
+        for (auto v : h_sorted) {
+          std::cout << v << ", ";
+        }
+        std::cout << std::endl;
+      }
+      CHECK((*p_tree)[nidx].IsLeaf() || n_leaf == 1) << " nidx:" << nidx;
+    }
+
+    /**
+     * copy node pointer
+     */
+    // thrust::inclusive_scan(thrust::cuda::par(caching), counts_out.begin(),
+    //                        counts_out.begin() + n_leaf, dh::tbegin(d_node_ptr) + 1);
+    // auto const& h_node_ptr = row_indices.node_ptr.ConstHostVector();
+    // auto total = h_node_ptr.back();
+    // if (total != this->ridx_.Size()) {
+    //   std::cout << "Counts, n_leaf: " << n_leaf << std::endl;
+    //   for (size_t i = 0; i < n_leaf; ++i) {
+    //     std::cout << counts_out[i] << ", ";
+    //   }
+    //   std::cout << std::endl;
+
+    //   std::cout << "nodes, n_leaf: " << n_leaf << std::endl;
+    //   for (auto nidx : h_node_idx) {
+    //     std::cout << nidx << ", ";
+    //   }
+    //   std::cout << std::endl;
+    // }
+
+    // CHECK_EQ(total, this->ridx_.Size());
   }
 
   /**
diff --git a/tests/python-gpu/test_gpu_updaters.py b/tests/python-gpu/test_gpu_updaters.py
index a3427b566360..1d13cbd35e80 100644
--- a/tests/python-gpu/test_gpu_updaters.py
+++ b/tests/python-gpu/test_gpu_updaters.py
@@ -3,7 +3,7 @@
 import gc
 import pytest
 import xgboost as xgb
-from hypothesis import given, strategies, assume, settings, note
+from hypothesis import given, strategies, assume, settings, note, reproduce_failure, seed
 
 sys.path.append("tests/python")
 import testing as tm
@@ -11,8 +11,8 @@
 
 
 parameter_strategy = strategies.fixed_dictionaries({
-    'max_depth': strategies.integers(0, 11),
-    'max_leaves': strategies.integers(0, 256),
+    'max_depth': strategies.integers(0, 4),
+    'max_leaves': strategies.integers(0, 8),
     'max_bin': strategies.integers(2, 1024),
     'grow_policy': strategies.sampled_from(['lossguide', 'depthwise']),
     'single_precision_histogram': strategies.booleans(),
@@ -47,6 +47,7 @@ class TestGPUUpdaters:
 
     @given(parameter_strategy, strategies.integers(1, 20), tm.dataset_strategy)
     @settings(deadline=None, print_blob=True)
+    @seed(1234)
     def test_gpu_hist(self, param, num_rounds, dataset):
         param["tree_method"] = "gpu_hist"
         param = dataset.set_params(param)
diff --git a/tests/python/testing.py b/tests/python/testing.py
index ad893d3ef15f..53eda94f5713 100644
--- a/tests/python/testing.py
+++ b/tests/python/testing.py
@@ -324,28 +324,28 @@ def make_categorical(
 
 _unweighted_datasets_strategy = strategies.sampled_from(
     [
-        TestDataset(
-            "calif_housing", get_california_housing, "reg:squarederror", "rmse"
-        ),
+        # TestDataset(
+        #     "calif_housing", get_california_housing, "reg:squarederror", "rmse"
+        # ),
         TestDataset(
             "calif_housing", get_california_housing, "reg:absoluteerror", "mae"
         ),
-        TestDataset("digits", get_digits, "multi:softmax", "mlogloss"),
-        TestDataset("cancer", get_cancer, "binary:logistic", "logloss"),
-        TestDataset(
-            "mtreg",
-            lambda: datasets.make_regression(n_samples=128, n_targets=3),
-            "reg:squarederror",
-            "rmse",
-        ),
-        TestDataset("sparse", get_sparse, "reg:squarederror", "rmse"),
-        TestDataset("sparse", get_sparse, "reg:absoluteerror", "mae"),
-        TestDataset(
-            "empty",
-            lambda: (np.empty((0, 100)), np.empty(0)),
-            "reg:squarederror",
-            "rmse",
-        ),
+        # TestDataset("digits", get_digits, "multi:softmax", "mlogloss"),
+        # TestDataset("cancer", get_cancer, "binary:logistic", "logloss"),
+        # TestDataset(
+        #     "mtreg",
+        #     lambda: datasets.make_regression(n_samples=128, n_targets=3),
+        #     "reg:squarederror",
+        #     "rmse",
+        # ),
+        # TestDataset("sparse", get_sparse, "reg:squarederror", "rmse"),
+        # TestDataset("sparse", get_sparse, "reg:absoluteerror", "mae"),
+        # TestDataset(
+        #     "empty",
+        #     lambda: (np.empty((0, 100)), np.empty(0)),
+        #     "reg:squarederror",
+        #     "rmse",
+        # ),
     ]
 )
 

From f953ab3f7de6db51f88b6f40ef4e246362c6c8b6 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 14:52:51 +0800
Subject: [PATCH 069/124] Fix.

---
 python-package/xgboost/training.py    |  8 +++---
 src/learner.cc                        |  1 +
 src/tree/gpu_hist/row_partitioner.cuh | 38 ++++++++++++++-------------
 3 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/python-package/xgboost/training.py b/python-package/xgboost/training.py
index 016a83a701f9..7b2a8d96ee03 100644
--- a/python-package/xgboost/training.py
+++ b/python-package/xgboost/training.py
@@ -190,10 +190,10 @@ def filename(pre, ext):
         try:
             bst.update(dtrain, i, obj)
         except XGBoostError:
-            import pickle
-            with open(filename("model", "pkl"), "bw") as fd:
-                pickle.dump(bst, fd)
-            dtrain.save_binary(filename("Xy", "dmatrix"))
+            # import pickle
+            # with open(filename("model", "pkl"), "bw") as fd:
+            #     pickle.dump(bst, fd)
+            # dtrain.save_binary(filename("Xy", "dmatrix"))
             raise
         if cb_container.after_iteration(bst, i, dtrain, evals):
             break
diff --git a/src/learner.cc b/src/learner.cc
index 2aed2cde00f3..72d0561e563d 100644
--- a/src/learner.cc
+++ b/src/learner.cc
@@ -423,6 +423,7 @@ class LearnerConfiguration : public Learner {
       obj_.reset(ObjFunction::Create(tparam_.objective, &generic_parameters_));
     }
     obj_->LoadConfig(objective_fn);
+    learner_model_param_.task = obj_->Task();
 
     tparam_.booster = get<String>(gradient_booster["name"]);
     if (!gbm_) {
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 9bebdde6e610..fb77025ba338 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -177,7 +177,7 @@ class RowPartitioner {
                         Sampledp sampledp) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
-    if (!task.zero_hess) {
+    if (!task.UpdateTreeLeaf()) {
       dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
         auto position = d_position[idx];
         RowIndexT ridx = d_ridx[idx];
@@ -240,7 +240,7 @@ class RowPartitioner {
         reinterpret_cast<bst_node_t*>(pinned.subspan(sizeof(size_t), sizeof(bst_node_t)).data());
     dh::safe_cuda(cudaMemcpyAsync(h_num_runs, d_num_runs_out.data(), sizeof(size_t),
                                   cudaMemcpyDeviceToHost, streams_[0]));
-    dh::safe_cuda(cudaMemcpyAsync(h_first_unique, d_unique_out.data(), sizeof(size_t),
+    dh::safe_cuda(cudaMemcpyAsync(h_first_unique, d_unique_out.data(), sizeof(bst_node_t),
                                   cudaMemcpyDeviceToHost, streams_[0]));
 
     /**
@@ -279,8 +279,9 @@ class RowPartitioner {
     if (*h_first_unique == kIgnoredTreePosition){
       *h_num_runs -= 1;  // sampled.
     }
+    CHECK_GT(*h_num_runs, 0);
     // shrink to omit the `kIgnoredTreePosition`.
-    row_indices.node_ptr.Resize(*h_num_runs);
+    row_indices.node_ptr.Resize(*h_num_runs + 1);
     row_indices.node_idx.Resize(*h_num_runs);
 
     // std::cout << "n_leaf: " << n_leaf << std::endl;
@@ -294,21 +295,22 @@ class RowPartitioner {
     // std::cout << std::endl;
 
     auto const& h_node_idx = row_indices.node_idx.ConstHostVector();
-    for (size_t i = 0; i < n_leaf; ++i) {
-      auto nidx = unique_out[i];
-      // std::cout << "nidx:" << nidx << std::endl;
-      if (!((*p_tree)[nidx].IsLeaf()) && n_leaf != 1) {
-        std::cout << __LINE__ << " sorted position" << std::endl;
-        std::vector<bst_node_t> h_sorted(position_.Size());
-        dh::CopyDeviceSpanToVector(
-            &h_sorted, common::Span<bst_node_t const>{sorted_position, position_.Size()});
-        for (auto v : h_sorted) {
-          std::cout << v << ", ";
-        }
-        std::cout << std::endl;
-      }
-      CHECK((*p_tree)[nidx].IsLeaf() || n_leaf == 1) << " nidx:" << nidx;
-    }
+    // std::cout << "h_node_idx.size():" << h_node_idx.size() << std::endl;
+    // for (size_t i = 0; i < h_node_idx.size(); ++i) {
+    //   auto nidx = h_node_idx[i];
+    //   // std::cout << "nidx:" << nidx << std::endl;
+    //   // if (!((*p_tree)[nidx].IsLeaf()) && n_leaf != 1) {
+    //   //   std::cout << __LINE__ << " sorted position" << std::endl;
+    //   //   std::vector<bst_node_t> h_sorted(position_.Size());
+    //   //   dh::CopyDeviceSpanToVector(
+    //   //       &h_sorted, common::Span<bst_node_t const>{sorted_position, position_.Size()});
+    //   //   for (auto v : h_sorted) {
+    //   //     std::cout << v << ", ";
+    //   //   }
+    //   //   std::cout << std::endl;
+    //   // }
+    //   CHECK((*p_tree)[nidx].IsLeaf() || n_leaf == 1) << " nidx:" << nidx;
+    // }
 
     /**
      * copy node pointer

From 1ce049ccf6d016029b67aae00d6bcef313792130 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 15:46:29 +0800
Subject: [PATCH 070/124] Fix.

---
 src/tree/gpu_hist/row_partitioner.cuh | 7 +++++--
 tests/python-gpu/test_gpu_updaters.py | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index fb77025ba338..6828a5d96751 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -267,10 +267,13 @@ class RowPartitioner {
         // smallest value and is sorted to the left.
         // d_unique_out.size() == n_leaf + 1.
         d_node_idx[i] = d_unique_out[i + 1];
-        d_node_ptr[i] = d_counts_out[i + 1];
+        d_node_ptr[i + 1] = d_counts_out[i + 1];
       } else {
         d_node_idx[i] = d_unique_out[i];
-        d_node_ptr[i] = d_counts_out[i];
+        d_node_ptr[i + 1] = d_counts_out[i];
+      }
+      if (i == 0) {
+        d_node_ptr[i] = 0;
       }
     });
     thrust::inclusive_scan(thrust::cuda::par(caching), dh::tbegin(d_node_ptr), dh::tend(d_node_ptr),
diff --git a/tests/python-gpu/test_gpu_updaters.py b/tests/python-gpu/test_gpu_updaters.py
index 1d13cbd35e80..a15c9798e328 100644
--- a/tests/python-gpu/test_gpu_updaters.py
+++ b/tests/python-gpu/test_gpu_updaters.py
@@ -47,11 +47,11 @@ class TestGPUUpdaters:
 
     @given(parameter_strategy, strategies.integers(1, 20), tm.dataset_strategy)
     @settings(deadline=None, print_blob=True)
-    @seed(1234)
     def test_gpu_hist(self, param, num_rounds, dataset):
         param["tree_method"] = "gpu_hist"
         param = dataset.set_params(param)
         result = train_result(param, dataset.get_dmat(), num_rounds)
+
         note(result)
         assert tm.non_increasing(result["train"][dataset.metric])
 

From d9d51e4211a4920a5d719d666489595332efbdc8 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 15:49:33 +0800
Subject: [PATCH 071/124] Cleanup.

---
 python-package/xgboost/training.py    | 18 +-----------
 tests/python-gpu/test_gpu_updaters.py |  4 +--
 tests/python/testing.py               | 40 +++++++++++++--------------
 3 files changed, 23 insertions(+), 39 deletions(-)

diff --git a/python-package/xgboost/training.py b/python-package/xgboost/training.py
index 7b2a8d96ee03..38567b6bf949 100644
--- a/python-package/xgboost/training.py
+++ b/python-package/xgboost/training.py
@@ -175,26 +175,10 @@ def train(
 
     bst = cb_container.before_training(bst)
 
-    def filename(pre, ext):
-        i = 0
-        fn = f"{pre}-" + str(i) + f".{ext}"
-        while os.path.exists(fn):
-            i += 1
-            fn = f"{pre}-" + str(i) + f".{ext}"
-        return fn
-
     for i in range(start_iteration, num_boost_round):
-        print("i: ", i)
         if cb_container.before_iteration(bst, i, dtrain, evals):
             break
-        try:
-            bst.update(dtrain, i, obj)
-        except XGBoostError:
-            # import pickle
-            # with open(filename("model", "pkl"), "bw") as fd:
-            #     pickle.dump(bst, fd)
-            # dtrain.save_binary(filename("Xy", "dmatrix"))
-            raise
+        bst.update(dtrain, i, obj)
         if cb_container.after_iteration(bst, i, dtrain, evals):
             break
 
diff --git a/tests/python-gpu/test_gpu_updaters.py b/tests/python-gpu/test_gpu_updaters.py
index a15c9798e328..9eb68a2d10a1 100644
--- a/tests/python-gpu/test_gpu_updaters.py
+++ b/tests/python-gpu/test_gpu_updaters.py
@@ -11,8 +11,8 @@
 
 
 parameter_strategy = strategies.fixed_dictionaries({
-    'max_depth': strategies.integers(0, 4),
-    'max_leaves': strategies.integers(0, 8),
+    'max_depth': strategies.integers(0, 11),
+    'max_leaves': strategies.integers(0, 256),
     'max_bin': strategies.integers(2, 1024),
     'grow_policy': strategies.sampled_from(['lossguide', 'depthwise']),
     'single_precision_histogram': strategies.booleans(),
diff --git a/tests/python/testing.py b/tests/python/testing.py
index 53eda94f5713..8633e4caa52d 100644
--- a/tests/python/testing.py
+++ b/tests/python/testing.py
@@ -324,28 +324,28 @@ def make_categorical(
 
 _unweighted_datasets_strategy = strategies.sampled_from(
     [
-        # TestDataset(
-        #     "calif_housing", get_california_housing, "reg:squarederror", "rmse"
-        # ),
         TestDataset(
-            "calif_housing", get_california_housing, "reg:absoluteerror", "mae"
+            "calif_housing", get_california_housing, "reg:squarederror", "rmse"
+        ),
+        TestDataset(
+            "calif_housing-l1", get_california_housing, "reg:absoluteerror", "mae"
+        ),
+        TestDataset("digits", get_digits, "multi:softmax", "mlogloss"),
+        TestDataset("cancer", get_cancer, "binary:logistic", "logloss"),
+        TestDataset(
+            "mtreg",
+            lambda: datasets.make_regression(n_samples=128, n_targets=3),
+            "reg:squarederror",
+            "rmse",
+        ),
+        TestDataset("sparse", get_sparse, "reg:squarederror", "rmse"),
+        TestDataset("sparse-l1", get_sparse, "reg:absoluteerror", "mae"),
+        TestDataset(
+            "empty",
+            lambda: (np.empty((0, 100)), np.empty(0)),
+            "reg:squarederror",
+            "rmse",
         ),
-        # TestDataset("digits", get_digits, "multi:softmax", "mlogloss"),
-        # TestDataset("cancer", get_cancer, "binary:logistic", "logloss"),
-        # TestDataset(
-        #     "mtreg",
-        #     lambda: datasets.make_regression(n_samples=128, n_targets=3),
-        #     "reg:squarederror",
-        #     "rmse",
-        # ),
-        # TestDataset("sparse", get_sparse, "reg:squarederror", "rmse"),
-        # TestDataset("sparse", get_sparse, "reg:absoluteerror", "mae"),
-        # TestDataset(
-        #     "empty",
-        #     lambda: (np.empty((0, 100)), np.empty(0)),
-        #     "reg:squarederror",
-        #     "rmse",
-        # ),
     ]
 )
 

From cdf93e2cab98970ff5ab7c6fe89e139e4befacbb Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 15:59:41 +0800
Subject: [PATCH 072/124] Cleanups.

---
 include/xgboost/linalg.h              |  1 -
 include/xgboost/task.h                |  1 -
 include/xgboost/tree_model.h          |  1 -
 src/common/device_helpers.cuh         |  6 +--
 src/common/partition_builder.h        |  1 -
 src/learner.cc                        | 36 +---------------
 src/objective/regression_obj.cc       |  8 ----
 src/objective/regression_obj.cu       |  1 -
 src/tree/gpu_hist/row_partitioner.cuh | 60 ---------------------------
 9 files changed, 4 insertions(+), 111 deletions(-)

diff --git a/include/xgboost/linalg.h b/include/xgboost/linalg.h
index 47082b42d8f1..015121560039 100644
--- a/include/xgboost/linalg.h
+++ b/include/xgboost/linalg.h
@@ -14,7 +14,6 @@
 
 #include <algorithm>
 #include <cassert>
-#include <cstddef>
 #include <limits>
 #include <string>
 #include <type_traits>
diff --git a/include/xgboost/task.h b/include/xgboost/task.h
index 22a65f2dad4b..739207a309d8 100644
--- a/include/xgboost/task.h
+++ b/include/xgboost/task.h
@@ -7,7 +7,6 @@
 #include <xgboost/base.h>
 
 #include <cinttypes>
-#include <cstdint>
 
 namespace xgboost {
 /*!
diff --git a/include/xgboost/tree_model.h b/include/xgboost/tree_model.h
index 3cebfa1d9095..789d336f59fb 100644
--- a/include/xgboost/tree_model.h
+++ b/include/xgboost/tree_model.h
@@ -16,7 +16,6 @@
 #include <xgboost/feature_map.h>
 #include <xgboost/model.h>
 
-#include <cstddef>
 #include <limits>
 #include <vector>
 #include <string>
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
index 71e0d851e7e0..9adf866fece9 100644
--- a/src/common/device_helpers.cuh
+++ b/src/common/device_helpers.cuh
@@ -1410,7 +1410,7 @@ void InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op,
 }
 
 template <typename InIt, typename OutIt, typename Predicate>
-OutIt CopyIf(InIt in_first, InIt in_second, OutIt out_first, Predicate pred) {
+void CopyIf(InIt in_first, InIt in_second, OutIt out_first, Predicate pred) {
   // We loop over batches because thrust::copy_if can't deal with sizes > 2^31
   // See thrust issue #1302, XGBoost #6822
   size_t constexpr kMaxCopySize = std::numeric_limits<int>::max() / 2;
@@ -1419,9 +1419,9 @@ OutIt CopyIf(InIt in_first, InIt in_second, OutIt out_first, Predicate pred) {
   for (size_t offset = 0; offset < length; offset += kMaxCopySize) {
     auto begin_input = in_first + offset;
     auto end_input = in_first + std::min(offset + kMaxCopySize, length);
-    out_first = thrust::copy_if(thrust::cuda::par(alloc), begin_input, end_input, out_first, pred);
+    out_first = thrust::copy_if(thrust::cuda::par(alloc), begin_input,
+                                end_input, out_first, pred);
   }
-  return out_first;
 }
 
 template <typename InputIteratorT, typename OutputIteratorT, typename OffsetT>
diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index 20173d15b11d..44f6a0ba464b 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -10,7 +10,6 @@
 #include <xgboost/data.h>
 
 #include <algorithm>
-#include <cstddef>
 #include <memory>
 #include <utility>
 #include <vector>
diff --git a/src/learner.cc b/src/learner.cc
index 72d0561e563d..c58656b9981b 100644
--- a/src/learner.cc
+++ b/src/learner.cc
@@ -7,11 +7,8 @@
 #include <dmlc/io.h>
 #include <dmlc/parameter.h>
 #include <dmlc/thread_local.h>
-#include <dmlc/filesystem.h>
 
 #include <atomic>
-#include <cstddef>
-#include <fstream>
 #include <mutex>
 #include <algorithm>
 #include <iomanip>
@@ -24,7 +21,6 @@
 #include <vector>
 
 #include "dmlc/any.h"
-#include "dmlc/logging.h"
 #include "xgboost/base.h"
 #include "xgboost/c_api.h"
 #include "xgboost/data.h"
@@ -1172,38 +1168,8 @@ class LearnerImpl : public LearnerIO {
     obj_->GetGradient(predt.predictions, train->Info(), iter, &gpair_);
     monitor_.Stop("GetGradient");
     TrainingObserver::Instance().Observe(gpair_, "Gradients");
-    std::stringstream ss;
-    ss << common::GlobalRandom();
-    auto str = ss.str();
 
-    auto FileExists = [](const std::string& filename) {
-      struct stat st;
-      return stat(filename.c_str(), &st) == 0;
-    };
-    auto __attribute__((unused)) SaveRng = [&]() {
-      size_t i = 0;
-      std::string filename{"rng-" + std::to_string(i)};
-      while (FileExists(filename)) {
-        ++i;
-        filename = "rng-" + std::to_string(i);
-      }
-      std::ofstream fout(filename);
-      fout << str;
-    };
-    auto __attribute__((unused)) LoadRng = []() {
-      std::string filename {"rng-1"};
-      std::ifstream fin(filename);
-      std::string str;
-      fin >> common::GlobalRandom();
-    };
-
-    LoadRng();
-    try {
-      gbm_->DoBoost(train.get(), &gpair_, &predt, obj_.get());
-    } catch (dmlc::Error const& e) {
-      // SaveRng();
-      throw e;
-    }
+    gbm_->DoBoost(train.get(), &gpair_, &predt);
     monitor_.Stop("UpdateOneIter");
   }
 
diff --git a/src/objective/regression_obj.cc b/src/objective/regression_obj.cc
index 0a80064850fd..663989fbd5c3 100644
--- a/src/objective/regression_obj.cc
+++ b/src/objective/regression_obj.cc
@@ -5,14 +5,6 @@
 // Dummy file to keep the CUDA conditional compile trick.
 
 #include <dmlc/registry.h>
-
-#include "../common/linalg_op.h"
-
-#include "rabit/rabit.h"
-#include "xgboost/data.h"
-#include "xgboost/objective.h"
-#include "xgboost/tree_model.h"
-
 namespace xgboost {
 namespace obj {
 
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 98944cc1b7c6..e01960ccb8c8 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -11,7 +11,6 @@
 #include <xgboost/tree_model.h>
 
 #include <cmath>
-#include <cstddef>
 #include <memory>
 #include <vector>
 
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 6828a5d96751..90d2fd87c14e 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -2,8 +2,6 @@
  * Copyright 2017-2019 XGBoost contributors
  */
 #pragma once
-#include <cstddef>
-#include <cstdint>
 #include <limits>
 #include <vector>
 #include "xgboost/base.h"
@@ -149,13 +147,6 @@ class RowPartitioner {
     CHECK_GE(left_count, 0);
     ridx_segments_.resize(std::max(static_cast<bst_node_t>(ridx_segments_.size()),
                                    std::max(left_nidx, right_nidx) + 1));
-    // std::cout << "nidx: " << nidx << ", left:" << left_nidx << ", right:" << right_nidx << " count: " << left_count << std::endl;
-    // if (left_count == segment.Size()) {
-    //   LOG(FATAL) << "empty right:" << left_count << ", nidx:" << nidx << std::endl;
-    // }
-    // if (left_count == 0) {
-    //   LOG(FATAL) << "empty left:" << segment.Size() << ", nidx:" << nidx << std::endl;
-    // }
     ridx_segments_[left_nidx] =
         Segment(segment.begin, segment.begin + left_count);
     ridx_segments_[right_nidx] =
@@ -286,57 +277,6 @@ class RowPartitioner {
     // shrink to omit the `kIgnoredTreePosition`.
     row_indices.node_ptr.Resize(*h_num_runs + 1);
     row_indices.node_idx.Resize(*h_num_runs);
-
-    // std::cout << "n_leaf: " << n_leaf << std::endl;
-    // auto const& self = *p_tree;
-    // p_tree->WalkTree([&self](bst_node_t nidx) {
-    //   if (self[nidx].IsLeaf()) {
-    //     std::cout << nidx << ", p:" << self[nidx].Parent() << ", ";
-    //   }
-    //   return true;
-    // });
-    // std::cout << std::endl;
-
-    auto const& h_node_idx = row_indices.node_idx.ConstHostVector();
-    // std::cout << "h_node_idx.size():" << h_node_idx.size() << std::endl;
-    // for (size_t i = 0; i < h_node_idx.size(); ++i) {
-    //   auto nidx = h_node_idx[i];
-    //   // std::cout << "nidx:" << nidx << std::endl;
-    //   // if (!((*p_tree)[nidx].IsLeaf()) && n_leaf != 1) {
-    //   //   std::cout << __LINE__ << " sorted position" << std::endl;
-    //   //   std::vector<bst_node_t> h_sorted(position_.Size());
-    //   //   dh::CopyDeviceSpanToVector(
-    //   //       &h_sorted, common::Span<bst_node_t const>{sorted_position, position_.Size()});
-    //   //   for (auto v : h_sorted) {
-    //   //     std::cout << v << ", ";
-    //   //   }
-    //   //   std::cout << std::endl;
-    //   // }
-    //   CHECK((*p_tree)[nidx].IsLeaf() || n_leaf == 1) << " nidx:" << nidx;
-    // }
-
-    /**
-     * copy node pointer
-     */
-    // thrust::inclusive_scan(thrust::cuda::par(caching), counts_out.begin(),
-    //                        counts_out.begin() + n_leaf, dh::tbegin(d_node_ptr) + 1);
-    // auto const& h_node_ptr = row_indices.node_ptr.ConstHostVector();
-    // auto total = h_node_ptr.back();
-    // if (total != this->ridx_.Size()) {
-    //   std::cout << "Counts, n_leaf: " << n_leaf << std::endl;
-    //   for (size_t i = 0; i < n_leaf; ++i) {
-    //     std::cout << counts_out[i] << ", ";
-    //   }
-    //   std::cout << std::endl;
-
-    //   std::cout << "nodes, n_leaf: " << n_leaf << std::endl;
-    //   for (auto nidx : h_node_idx) {
-    //     std::cout << nidx << ", ";
-    //   }
-    //   std::cout << std::endl;
-    // }
-
-    // CHECK_EQ(total, this->ridx_.Size());
   }
 
   /**

From 914c19a7c047acbd406fb9ffe711ff8318d18b6c Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 16:03:01 +0800
Subject: [PATCH 073/124] Cleanup.

---
 src/common/stats.cuh                            | 1 -
 src/tree/updater_approx.h                       | 1 -
 src/tree/updater_gpu_hist.cu                    | 4 ----
 tests/cpp/common/test_stats.cu                  | 1 -
 tests/cpp/tree/gpu_hist/test_row_partitioner.cu | 1 -
 tests/python/test_with_dask.py                  | 3 +--
 6 files changed, 1 insertion(+), 10 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index b01089a592a8..00b916fe0b1e 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -8,7 +8,6 @@
 #include <thrust/iterator/permutation_iterator.h>
 #include <thrust/sort.h>
 
-#include <cstddef>
 #include <iterator>  // std::distance
 
 #include "device_helpers.cuh"
diff --git a/src/tree/updater_approx.h b/src/tree/updater_approx.h
index b662c6c4e81f..1e8f38da4269 100644
--- a/src/tree/updater_approx.h
+++ b/src/tree/updater_approx.h
@@ -6,7 +6,6 @@
 #ifndef XGBOOST_TREE_UPDATER_APPROX_H_
 #define XGBOOST_TREE_UPDATER_APPROX_H_
 
-#include <cstddef>
 #include <limits>
 #include <utility>
 #include <vector>
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index e8e5902838bb..d38cdf92dacb 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -6,7 +6,6 @@
 #include <xgboost/tree_updater.h>
 #include <algorithm>
 #include <cmath>
-#include <cstddef>
 #include <memory>
 #include <limits>
 #include <utility>
@@ -355,7 +354,6 @@ struct GPUHistMakerDevice {
   }
 
   void UpdatePosition(int nidx, RegTree* p_tree) {
-    // std::cout << "UpdatePosition:" << nidx << std::endl;
     RegTree::Node split_node = (*p_tree)[nidx];
     auto split_type = p_tree->NodeSplitType(nidx);
     auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
@@ -384,7 +382,6 @@ struct GPUHistMakerDevice {
               new_position = split_node.RightChild();
             }
           }
-          // printf("ridx: %d, pos: %d\n", ridx, new_position);
           return new_position;
         });
   }
@@ -450,7 +447,6 @@ struct GPUHistMakerDevice {
         [=] __device__(size_t row_id, int position) {
           // What happens if user prune the tree?
           if (!d_matrix.IsInRange(row_id)) {
-            // printf("out\n");
             return RowPartitioner::kIgnoredTreePosition;
           }
           auto node = d_nodes[position];
diff --git a/tests/cpp/common/test_stats.cu b/tests/cpp/common/test_stats.cu
index fb28d8e4da98..eee92921d931 100644
--- a/tests/cpp/common/test_stats.cu
+++ b/tests/cpp/common/test_stats.cu
@@ -2,7 +2,6 @@
  * Copyright 2022 by XGBoost Contributors
  */
 #include <gtest/gtest.h>
-#include <cstddef>
 #include <utility>
 #include <vector>
 
diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index a203e75b7b5d..7e0de4dea9a0 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -3,7 +3,6 @@
  */
 #include <gtest/gtest.h>
 #include <algorithm>
-#include <cstddef>
 #include <vector>
 
 #include <thrust/device_vector.h>
diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index b6a8297dd91f..5f1015f11278 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -1292,8 +1292,7 @@ def minimum_bin():
         if minimum_bin() and is_stump():
             assert tm.non_increasing(history, tolerance=1e-3)
         else:
-            pass
-            # assert tm.non_increasing(history)
+            assert tm.non_increasing(history)
         # Make sure that it's decreasing
         assert history[-1] < history[0]
 

From 0a5525e4c0ca7ca33cc1eda7530bb9530ee59e38 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 16:27:55 +0800
Subject: [PATCH 074/124] Fix test.

---
 src/learner.cc                                  |  2 +-
 src/tree/gpu_hist/row_partitioner.cuh           | 13 ++++++++-----
 tests/cpp/tree/gpu_hist/test_row_partitioner.cu | 14 ++++++--------
 tests/python-gpu/test_gpu_with_dask.py          |  3 +--
 4 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/src/learner.cc b/src/learner.cc
index c58656b9981b..1fc987d65427 100644
--- a/src/learner.cc
+++ b/src/learner.cc
@@ -1169,7 +1169,7 @@ class LearnerImpl : public LearnerIO {
     monitor_.Stop("GetGradient");
     TrainingObserver::Instance().Observe(gpair_, "Gradients");
 
-    gbm_->DoBoost(train.get(), &gpair_, &predt);
+    gbm_->DoBoost(train.get(), &gpair_, &predt, obj_.get());
     monitor_.Stop("UpdateOneIter");
   }
 
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 90d2fd87c14e..7e5752b361ad 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -212,7 +212,7 @@ class RowPartitioner {
     dh::caching_device_vector<size_t> counts_out(max_n_unique + 1, 0);
     auto d_counts_out = dh::ToSpan(counts_out).subspan(0, max_n_unique);
     auto d_num_runs_out = dh::ToSpan(counts_out).subspan(max_n_unique, 1);
-    dh::device_vector<bst_node_t> unique_out(max_n_unique, 0);
+    dh::caching_device_vector<bst_node_t> unique_out(max_n_unique, 0);
     auto d_unique_out = dh::ToSpan(unique_out);
 
     size_t nbytes;
@@ -259,18 +259,21 @@ class RowPartitioner {
         // d_unique_out.size() == n_leaf + 1.
         d_node_idx[i] = d_unique_out[i + 1];
         d_node_ptr[i + 1] = d_counts_out[i + 1];
+        if (i == 0) {
+          d_node_ptr[0] = d_counts_out[0];
+        }
       } else {
         d_node_idx[i] = d_unique_out[i];
         d_node_ptr[i + 1] = d_counts_out[i];
-      }
-      if (i == 0) {
-        d_node_ptr[i] = 0;
+        if (i == 0) {
+          d_node_ptr[0] = 0;
+        }
       }
     });
     thrust::inclusive_scan(thrust::cuda::par(caching), dh::tbegin(d_node_ptr), dh::tend(d_node_ptr),
                            dh::tbegin(d_node_ptr));
     dh::CUDAStreamView{streams_[0]}.Sync();
-    if (*h_first_unique == kIgnoredTreePosition){
+    if (*h_first_unique == kIgnoredTreePosition) {
       *h_num_runs -= 1;  // sampled.
     }
     CHECK_GT(*h_num_runs, 0);
diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index 7e0de4dea9a0..755ef6f42d55 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -158,9 +158,9 @@ void TestFinalise() {
 
   auto const& h_node_ptr = row_index.back().node_ptr.ConstHostVector();
   ASSERT_EQ(h_node_ptr.size(), 3);
-  ASSERT_EQ(h_node_ptr[0], 0);
-  ASSERT_EQ(h_node_ptr[1], 3);
-  ASSERT_EQ(h_node_ptr[2], 6);
+  ASSERT_EQ(h_node_ptr[0], 4);
+  ASSERT_EQ(h_node_ptr[1], 7);
+  ASSERT_EQ(h_node_ptr[2], kNumRows);
 
   auto const& h_node_idx = row_index.back().node_idx.ConstHostVector();
   ASSERT_EQ(h_node_idx.size(), 2);
@@ -168,11 +168,9 @@ void TestFinalise() {
   ASSERT_EQ(h_node_idx[1], 2);
 
   auto const& h_ridx = row_index.back().row_index.ConstHostVector();
-  for (size_t i = h_node_ptr[0]; i < h_node_ptr[1]; ++i) {
-    ASSERT_EQ(h_ridx[i] % 2, 0);
-  }
-  for (size_t i = h_node_ptr[1]; i < h_node_ptr[2]; ++i) {
-    ASSERT_EQ(h_ridx[i] % 2, 1);
+  std::vector<size_t> sol{0, 3, 6, 9, 2, 4, 8, 1, 5, 7};
+  for (size_t i = 0; i < h_ridx.size(); ++i) {
+    ASSERT_EQ(h_ridx[i], sol[i]);
   }
 }
 
diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
index 1f0339e913ec..645ad01fa2ae 100644
--- a/tests/python-gpu/test_gpu_with_dask.py
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -305,8 +305,7 @@ def test_dask_classifier(
 
     def test_empty_dmatrix(self, local_cuda_cluster: LocalCUDACluster) -> None:
         with Client(local_cuda_cluster) as client:
-            parameters = {'tree_method': 'gpu_hist',
-                          'debug_synchronize': True}
+            parameters = {'tree_method': 'gpu_hist', 'debug_synchronize': True}
             run_empty_dmatrix_reg(client, parameters)
             run_empty_dmatrix_cls(client, parameters)
 

From 8af41ad197081638e51a37431844b9f1bdfb015c Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 17:26:41 +0800
Subject: [PATCH 075/124] Fix test.

---
 src/objective/regression_obj.cu            | 2 --
 tests/cpp/objective/test_regression_obj.cc | 1 +
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index e01960ccb8c8..b3f6a903e1db 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -714,8 +714,6 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
 
   part.row_index.SetDevice(ctx->gpu_id);
   auto d_row_index = part.row_index.ConstDeviceSpan();
-  auto const& h_node_ptr = part.node_ptr.ConstHostVector();
-  CHECK_LE(h_node_ptr.back(), info.num_row_);
   part.node_ptr.SetDevice(ctx->gpu_id);
   auto seg_beg = part.node_ptr.ConstDeviceSpan().data();
   auto seg_end = seg_beg + part.node_ptr.Size();
diff --git a/tests/cpp/objective/test_regression_obj.cc b/tests/cpp/objective/test_regression_obj.cc
index 2b36c4d087da..f561c939345d 100644
--- a/tests/cpp/objective/test_regression_obj.cc
+++ b/tests/cpp/objective/test_regression_obj.cc
@@ -389,6 +389,7 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteError)) {
   std::vector<float> labels{0.f, 3.f, 2.f, 5.f, 4.f, 7.f};
   info.labels.Reshape(6, 1);
   info.labels.Data()->HostVector() = labels;
+  info.num_row_ = labels.size();
   HostDeviceVector<float> predt{1.f, 2.f, 3.f, 4.f, 5.f, 6.f};
   info.weights_.HostVector() = {1.f, 1.f, 1.f, 1.f, 1.f, 1.f};
 

From aca330408dc536f5172a3731244c322be4ee8be0 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 18:50:25 +0800
Subject: [PATCH 076/124] Fill missing leaf.

---
 src/objective/regression_obj.cu               |  2 +-
 src/tree/gpu_hist/row_partitioner.cuh         | 52 ++++++++++++++++---
 src/tree/updater_gpu_hist.cu                  | 11 ++--
 .../cpp/tree/gpu_hist/test_row_partitioner.cu | 39 +++++++++++---
 tests/python-gpu/test_gpu_with_dask.py        | 10 +++-
 tests/python/test_with_dask.py                |  5 ++
 6 files changed, 99 insertions(+), 20 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index b3f6a903e1db..2adde6886ea2 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -676,7 +676,7 @@ void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_
   auto& tree = *p_tree;
   auto& quantiles = *p_quantiles;
 
-  size_t n_leaf{quantiles.size()};
+  size_t n_leaf{row_index.node_idx.Size()};
   rabit::Allreduce<rabit::op::Max>(&n_leaf, 1);
   CHECK(quantiles.empty() || quantiles.size() == n_leaf);
   if (quantiles.empty()) {
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 7e5752b361ad..5fa0318fd0bf 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -12,6 +12,27 @@
 
 namespace xgboost {
 namespace tree {
+namespace detail {
+inline void FillMissingLeaf(std::vector<bst_node_t> const& missing, RowIndexCache* p_row_indices) {
+  auto& row_indices = *p_row_indices;
+  auto& h_node_idx = row_indices.node_idx.HostVector();
+  auto& h_node_ptr = row_indices.node_ptr.HostVector();
+
+  for (auto leaf : missing) {
+    if (std::binary_search(h_node_idx.cbegin(), h_node_idx.cend(), leaf)) {
+      continue;
+    }
+    auto it = std::upper_bound(h_node_idx.cbegin(), h_node_idx.cend(), leaf);
+    auto pos = it - h_node_idx.cbegin();
+    h_node_idx.insert(h_node_idx.cbegin() + pos, leaf);
+    h_node_ptr.insert(h_node_ptr.cbegin() + pos, h_node_ptr[pos]);
+  }
+
+  // push to device.
+  row_indices.node_idx.ConstDevicePointer();
+  row_indices.node_ptr.ConstDevicePointer();
+}
+}  // namespace detail
 
 /*! \brief Count how many rows are assigned to left node. */
 __forceinline__ __device__ void AtomicIncrement(int64_t* d_count, bool increment) {
@@ -163,7 +184,7 @@ class RowPartitioner {
    * argument and return the new position for this training instance.
    */
   template <typename FinalisePositionOpT, typename Sampledp>
-  void FinalisePosition(Context const* ctx, RegTree const* p_tree, ObjInfo task,
+  void FinalisePosition(Context const* ctx, RegTree const* p_tree, size_t n_leaf, ObjInfo task,
                         std::vector<RowIndexCache>* p_out_row_indices, FinalisePositionOpT op,
                         Sampledp sampledp) {
     auto d_position = position_.Current();
@@ -204,8 +225,6 @@ class RowPartitioner {
     thrust::stable_sort_by_key(thrust::cuda::par(alloc), sorted_position,
                                sorted_position + n_samples, row_indices.row_index.DevicePointer());
 
-    // fixme: we know how many leaves there are
-    size_t n_leaf = p_tree->GetNumLeaves();
     // +1 for subsample, which is set to an unique value in above kernel.
     size_t max_n_unique = n_leaf + 1;
 
@@ -277,9 +296,30 @@ class RowPartitioner {
       *h_num_runs -= 1;  // sampled.
     }
     CHECK_GT(*h_num_runs, 0);
-    // shrink to omit the `kIgnoredTreePosition`.
-    row_indices.node_ptr.Resize(*h_num_runs + 1);
-    row_indices.node_idx.Resize(*h_num_runs);
+    CHECK_LE(*h_num_runs, n_leaf);
+
+    if (*h_num_runs < n_leaf) {
+      // empty leaf, have to fill in all missing leaves
+      auto const& tree = *p_tree;
+      // shrink to omit the `kIgnoredTreePosition`.
+      row_indices.node_ptr.Resize(*h_num_runs + 1);
+      row_indices.node_idx.Resize(*h_num_runs);
+
+      std::vector<bst_node_t> leaves;
+      tree.WalkTree([&](bst_node_t nidx) {
+        if (tree[nidx].IsLeaf()) {
+          leaves.push_back(nidx);
+        }
+        return true;
+      });
+      CHECK_EQ(leaves.size(), n_leaf);
+      // Fill all the leaves that don't have any sample. This is hacky and inefficien. An
+      // alternative is to leave the objective to handle missing leaf, which is more messy
+      // as we need to take other distributed workers into account.
+      detail::FillMissingLeaf(leaves, &row_indices);
+    }
+    CHECK_EQ(row_indices.node_idx.Size(), n_leaf);
+    CHECK_EQ(row_indices.node_ptr.Size(), n_leaf + 1);
   }
 
   /**
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index d38cdf92dacb..9ba3ab36cb19 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -389,7 +389,7 @@ struct GPUHistMakerDevice {
   // After tree update is finished, update the position of all training
   // instances to their final leaf. This information is used later to update the
   // prediction cache
-  void FinalisePosition(RegTree const* p_tree, DMatrix* p_fmat, ObjInfo task,
+  void FinalisePosition(RegTree const* p_tree, size_t n_leaf, DMatrix* p_fmat, ObjInfo task,
                         std::vector<RowIndexCache>* p_out_row_indices) {
     dh::TemporaryArray<RegTree::Node> d_nodes(p_tree->GetNodes().size());
     dh::safe_cuda(cudaMemcpyAsync(d_nodes.data().get(), p_tree->GetNodes().data(),
@@ -422,12 +422,12 @@ struct GPUHistMakerDevice {
     if (page->n_rows == p_fmat->Info().num_row_) {
       FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
                              dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
-                             p_out_row_indices);
+                             n_leaf, p_out_row_indices);
     } else {
       for (auto const& batch : p_fmat->GetBatches<EllpackPage>(batch_param)) {
         FinalisePositionInPage(batch.Impl(), p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
                                dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
-                               p_out_row_indices);
+                               n_leaf, p_out_row_indices);
       }
     }
   }
@@ -439,11 +439,12 @@ struct GPUHistMakerDevice {
                               common::Span<uint32_t const> categories,
                               common::Span<RegTree::Segment> categories_segments,
                               ObjInfo task,
+                              size_t n_leaf,
                               std::vector<RowIndexCache>* p_out_row_indices) {
     auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
     auto d_gpair = this->gpair;
     row_partitioner->FinalisePosition(
-        ctx_, p_tree, task, p_out_row_indices,
+        ctx_, p_tree, n_leaf, task, p_out_row_indices,
         [=] __device__(size_t row_id, int position) {
           // What happens if user prune the tree?
           if (!d_matrix.IsInRange(row_id)) {
@@ -712,7 +713,7 @@ struct GPUHistMakerDevice {
     }
 
     monitor.Start("FinalisePosition");
-    this->FinalisePosition(p_tree, p_fmat, task, p_out_row_indices);
+    this->FinalisePosition(p_tree, num_leaves, p_fmat, task, p_out_row_indices);
     monitor.Stop("FinalisePosition");
   }
 };
diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index 755ef6f42d55..363e209519e5 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -119,7 +119,7 @@ void TestFinalise() {
   {
     RowPartitioner rp(0, kNumRows);
     rp.FinalisePosition(
-        &ctx, &tree, task, &row_index,
+        &ctx, &tree, tree.GetNumLeaves(), task, &row_index,
         [=] __device__(RowPartitioner::RowIndexT ridx, int position) { return 7; },
         [] XGBOOST_DEVICE(size_t idx) { return false; });
 
@@ -148,13 +148,11 @@ void TestFinalise() {
 
   RowPartitioner rp(0, kNumRows);
   rp.FinalisePosition(
-      &ctx, &tree, task, &row_index,
+      &ctx, &tree, tree.GetNumLeaves(), task, &row_index,
       [] __device__(RowPartitioner::RowIndexT ridx, bst_node_t position) {
         return ridx % 2 == 0 ? 1 : 2;
       },
-      [d_hess] __device__(size_t ridx) {
-        return d_hess[ridx] - 0.f == 0.f;
-      });
+      [d_hess] __device__(size_t ridx) { return d_hess[ridx] - 0.f == 0.f; });
 
   auto const& h_node_ptr = row_index.back().node_ptr.ConstHostVector();
   ASSERT_EQ(h_node_ptr.size(), 3);
@@ -174,7 +172,36 @@ void TestFinalise() {
   }
 }
 
-TEST(RowPartitioner, Finalise) { TestFinalise(); }
+void TestFillMissingLeaf() {
+  std::vector<bst_node_t> missing{1, 3};
+  Context ctx;
+  RowIndexCache row_index(&ctx, 10);
+  row_index.node_idx = {2, 4, 5};
+  row_index.node_ptr = {0, 4, 8, 16};
+
+  detail::FillMissingLeaf(missing, &row_index);
+
+  auto const& h_nidx = row_index.node_idx.HostVector();
+  auto const& h_nptr = row_index.node_ptr.HostVector();
+
+  ASSERT_EQ(h_nidx[0], missing[0]);
+  ASSERT_EQ(h_nidx[2], missing[1]);
+  ASSERT_EQ(h_nidx[1], 2);
+  ASSERT_EQ(h_nidx[3], 4);
+  ASSERT_EQ(h_nidx[4], 5);
+
+  ASSERT_EQ(h_nptr[0], 0);
+  ASSERT_EQ(h_nptr[1], 0);  // empty
+  ASSERT_EQ(h_nptr[2], 4);
+  ASSERT_EQ(h_nptr[3], 4);  // empty
+  ASSERT_EQ(h_nptr[4], 8);
+  ASSERT_EQ(h_nptr[5], 16);
+}
+
+TEST(RowPartitioner, Finalise) {
+  TestFillMissingLeaf();
+  TestFinalise();
+}
 
 void TestIncorrectRow() {
   RowPartitioner rp(0, 1);
diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
index 645ad01fa2ae..a13ff57ce245 100644
--- a/tests/python-gpu/test_gpu_with_dask.py
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -9,7 +9,7 @@
 import subprocess
 from collections import OrderedDict
 from inspect import signature
-from hypothesis import given, strategies, settings, note
+from hypothesis import given, strategies, settings, note, reproduce_failure
 from hypothesis._settings import duration
 from test_gpu_updaters import parameter_strategy
 
@@ -200,7 +200,12 @@ def run_gpu_hist(
         evals=[(m, "train")],
     )["history"]
     note(history)
-    assert tm.non_increasing(history["train"][dataset.metric])
+
+    # the approximate quantile for leaf can cause error on distributed training
+    if dataset.objective.endswith("-l1"):
+        assert tm.non_increasing(history["train"][dataset.metric], tolerance=1e-3)
+    else:
+        assert tm.non_increasing(history["train"][dataset.metric])
 
 
 @pytest.mark.skipif(**tm.no_cudf())
@@ -236,6 +241,7 @@ def test_dask_dataframe(self, local_cuda_cluster: LocalCUDACluster) -> None:
     @pytest.mark.parametrize(
         "local_cuda_cluster", [{"n_workers": 2}], indirect=["local_cuda_cluster"]
     )
+    # @reproduce_failure('6.14.8', b'AXicY2ZgYGBmZMAOGIESAAD+AAo=')
     def test_gpu_hist(
         self,
         params: Dict,
diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index 5f1015f11278..0e4cc1405615 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -1291,6 +1291,9 @@ def minimum_bin():
 
         if minimum_bin() and is_stump():
             assert tm.non_increasing(history, tolerance=1e-3)
+        elif dataset.objective.endswith("-l1"):
+            # the approximate quantile for leaf can cause error on distributed training
+            assert tm.non_increasing(history, tolerance=1e-3)
         else:
             assert tm.non_increasing(history)
         # Make sure that it's decreasing
@@ -1299,6 +1302,7 @@ def minimum_bin():
     @given(params=hist_parameter_strategy,
            dataset=tm.dataset_strategy)
     @settings(deadline=None, suppress_health_check=suppress, print_blob=True)
+    @reproduce_failure('6.27.2', b'AAAAAAAAAAEAAA==')
     def test_hist(
             self, params: Dict, dataset: tm.TestDataset, client: "Client"
     ) -> None:
@@ -1308,6 +1312,7 @@ def test_hist(
     @given(params=exact_parameter_strategy,
            dataset=tm.dataset_strategy)
     @settings(deadline=None, suppress_health_check=suppress, print_blob=True)
+    # @reproduce_failure('6.27.2', b'AXicY2JnNDpff6rLhjXq6iuBo0IcVjvut1lMgdlLZwVcYuBofTkvQOMDKgAyAcAGTcPCw==')
     # @reproduce_failure('6.36.1', b'AXicY2BkIBUwArUAAAB0AAQ=')
     def test_approx(
             self, client: "Client", params: Dict, dataset: tm.TestDataset

From 705d738ab7ac5d8daa18708bf5064a63842d902f Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 18:54:31 +0800
Subject: [PATCH 077/124] another sample.

---
 tests/python-gpu/test_gpu_with_dask.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
index a13ff57ce245..2ea4157c9045 100644
--- a/tests/python-gpu/test_gpu_with_dask.py
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -241,6 +241,7 @@ def test_dask_dataframe(self, local_cuda_cluster: LocalCUDACluster) -> None:
     @pytest.mark.parametrize(
         "local_cuda_cluster", [{"n_workers": 2}], indirect=["local_cuda_cluster"]
     )
+    @reproduce_failure('6.14.8', b'AXicY2ZgYGBiwAGYGBkYAADJAAk=')
     # @reproduce_failure('6.14.8', b'AXicY2ZgYGBmZMAOGIESAAD+AAo=')
     def test_gpu_hist(
         self,

From a23361f00873746a8d348b7664245e3fe504bbb7 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 19:04:50 +0800
Subject: [PATCH 078/124] Fix external memory.

---
 src/objective/regression_obj.cu | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 2adde6886ea2..cf0d56986a1e 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -686,7 +686,7 @@ void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_
   rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
   auto world = rabit::GetWorldSize();
   std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
-                 [&](float q) { return q / world; });
+                 [&](float q) { return q / static_cast<double>(world); });
 
   // fixme: verify this is correct for external memory
   auto const& h_node_idx = row_index.node_idx.HostVector();
@@ -790,6 +790,8 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
       quantiles[i] += results[i];
     }
   }
+  std::transform(quantiles.cbegin(), quantiles.cend(), quantiles.begin(),
+                 [&](float q) { return q / static_cast<float>(row_index.size()); });
 
   UpdateLeafValues(&quantiles, row_index.front(), p_tree);
 }

From 4948886b156e596ac4ac5a49fda4279dd3038217 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 19:07:46 +0800
Subject: [PATCH 079/124] Note.

---
 src/tree/updater_gpu_hist.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 9ba3ab36cb19..4b5f00b57904 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -417,7 +417,7 @@ struct GPUHistMakerDevice {
     }
     if (task.UpdateTreeLeaf() && !p_fmat->SingleColBlock() && param.subsample != 1.0) {
       // see comment in the `FinalisePositionInPage`.
-      LOG(FATAL) << "Current objective function can not be used with external memory";
+      LOG(FATAL) << "Current objective function can not be used with subsampled external memory.";
     }
     if (page->n_rows == p_fmat->Info().num_row_) {
       FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),

From 30783fde108510b02ff3810a50e0bf7be62b68c8 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 19:39:23 +0800
Subject: [PATCH 080/124] Return NAN.

---
 src/common/stats.cuh                            | 4 ++--
 tests/cpp/tree/gpu_hist/test_row_partitioner.cu | 2 ++
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index 00b916fe0b1e..fe91c8dce87e 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -84,7 +84,7 @@ void SegmentedQuantile(Context const* ctx, double alpha, SegIt seg_begin, SegIt
     size_t begin = seg_begin[seg_idx];
     auto n = static_cast<double>(seg_begin[seg_idx + 1] - begin);
     if (n == 0) {
-      d_results[i] = 0;
+      d_results[i] = std::numeric_limits<float>::quiet_NaN();
       return;
     }
 
@@ -138,7 +138,7 @@ void SegmentedWeightedQuantile(Context const* ctx, double alpha, SegIt seg_beg,
     size_t begin = seg_beg[seg_idx];
     auto n = static_cast<double>(seg_beg[seg_idx + 1] - begin);
     if (n == 0) {
-      d_results[i] = 0;
+      d_results[i] = std::numeric_limits<float>::quiet_NaN();
       return;
     }
     auto leaf_cdf = d_weight_cdf.subspan(begin, static_cast<size_t>(n));
diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index 363e209519e5..96f464a6126a 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -178,6 +178,8 @@ void TestFillMissingLeaf() {
   RowIndexCache row_index(&ctx, 10);
   row_index.node_idx = {2, 4, 5};
   row_index.node_ptr = {0, 4, 8, 16};
+  row_index.node_idx.SetDevice(0);
+  row_index.node_ptr.SetDevice(0);
 
   detail::FillMissingLeaf(missing, &row_index);
 

From 81014b254c5728532640b7fe6ded2e12149dceb8 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 21:37:28 +0800
Subject: [PATCH 081/124] Cleanup.

---
 cmake/Utils.cmake               |  6 ++++++
 tests/cpp/tree/test_gpu_hist.cu | 14 +++-----------
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/cmake/Utils.cmake b/cmake/Utils.cmake
index 963c494ccf26..687d8e2ebc58 100644
--- a/cmake/Utils.cmake
+++ b/cmake/Utils.cmake
@@ -152,6 +152,12 @@ function(xgboost_set_cuda_flags target)
       $<$<COMPILE_LANGUAGE:CUDA>:-lineinfo>)
   endif (USE_DEVICE_DEBUG)
 
+  if (FORCE_COLORED_OUTPUT AND (CMAKE_GENERATOR STREQUAL "Ninja") AND
+      ((CMAKE_CXX_COMPILER_ID STREQUAL "GNU") OR
+        (CMAKE_CXX_COMPILER_ID STREQUAL "Clang")))
+    target_compile_options(${target} PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fdiagnostics-color=always>)
+  endif()
+
   if (USE_NVTX)
     enable_nvtx(${target})
   endif (USE_NVTX)
diff --git a/tests/cpp/tree/test_gpu_hist.cu b/tests/cpp/tree/test_gpu_hist.cu
index 20a1978b08bf..f1baa53eeead 100644
--- a/tests/cpp/tree/test_gpu_hist.cu
+++ b/tests/cpp/tree/test_gpu_hist.cu
@@ -23,14 +23,6 @@
 
 namespace xgboost {
 namespace tree {
-namespace {
-auto MakeCtx() {
-  Context ctx;
-  ctx.gpu_id = 0;
-  return ctx;
-}
-}  // anonymous namespace
-
 TEST(GpuHist, DeviceHistogram) {
   // Ensures that node allocates correctly after reaching `kStopGrowingSize`.
   dh::safe_cuda(cudaSetDevice(0));
@@ -89,7 +81,7 @@ void TestBuildHist(bool use_shared_memory_histograms) {
   param.Init(args);
   auto page = BuildEllpackPage(kNRows, kNCols);
   BatchParam batch_param{};
-  Context ctx{MakeCtx()};
+  Context ctx{CreateEmptyGenericParam(0)};
   GPUHistMakerDevice<GradientSumT> maker(&ctx, page.get(), {}, kNRows, param, kNCols, kNCols,
                                          batch_param);
   xgboost::SimpleLCG gen;
@@ -167,7 +159,7 @@ TEST(GpuHist, ApplySplit) {
   BatchParam bparam;
   bparam.gpu_id = 0;
   bparam.max_bin = 3;
-  Context ctx{MakeCtx()};
+  Context ctx{CreateEmptyGenericParam(0)};
 
   for (auto& ellpack : m->GetBatches<EllpackPage>(bparam)){
     auto impl = ellpack.Impl();
@@ -233,7 +225,7 @@ TEST(GpuHist, EvaluateRootSplit) {
   // Initialize GPUHistMakerDevice
   auto page = BuildEllpackPage(kNRows, kNCols);
   BatchParam batch_param{};
-  Context ctx{MakeCtx()};
+  Context ctx{CreateEmptyGenericParam(0)};
   GPUHistMakerDevice<GradientPairPrecise> maker(&ctx, page.get(), {}, kNRows, param, kNCols, kNCols,
                                                 batch_param);
   // Initialize GPUHistMakerDevice::node_sum_gradients

From 3ab84a26fa71981089ac97e447ac1ca01fca5d7d Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 22:41:40 +0800
Subject: [PATCH 082/124] Fix subsample.

---
 src/common/device_helpers.cuh              | 37 +++++++++++++++++++
 src/common/stats.cuh                       | 35 ++----------------
 src/objective/regression_obj.cu            | 15 ++++----
 tests/cpp/objective/test_regression_obj.cc | 43 +++++++++++++++++++++-
 4 files changed, 90 insertions(+), 40 deletions(-)

diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
index 9adf866fece9..b4da11b559d5 100644
--- a/src/common/device_helpers.cuh
+++ b/src/common/device_helpers.cuh
@@ -1537,6 +1537,43 @@ void SegmentedArgSort(xgboost::common::Span<U> values,
                             sorted_idx.size_bytes(), cudaMemcpyDeviceToDevice));
 }
 
+/**
+ * \brief Different from the above one, this one can handle cases where segment doesn't
+ *        start from 0, but as a result it uses comparison sort.
+ */
+template <typename SegIt, typename ValIt>
+void SegmentedArgSort(SegIt seg_begin, SegIt seg_end, ValIt val_begin, ValIt val_end,
+                      dh::device_vector<size_t> *p_sorted_idx) {
+  using Tup = thrust::tuple<int32_t, float>;
+  auto &sorted_idx = *p_sorted_idx;
+  size_t n = std::distance(val_begin, val_end);
+  sorted_idx.resize(n);
+  dh::Iota(dh::ToSpan(sorted_idx));
+  dh::device_vector<Tup> keys(sorted_idx.size());
+  auto key_it = dh::MakeTransformIterator<Tup>(thrust::make_counting_iterator(0ul),
+                                               [=] XGBOOST_DEVICE(size_t i) -> Tup {
+                                                 int32_t leaf_idx;
+                                                 if (i < *seg_begin) {
+                                                   leaf_idx = -1;
+                                                 } else {
+                                                   leaf_idx = dh::SegmentId(seg_begin, seg_end, i);
+                                                 }
+                                                 auto residue = val_begin[i];
+                                                 return thrust::make_tuple(leaf_idx, residue);
+                                               });
+  dh::XGBCachingDeviceAllocator<char> caching;
+  thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
+
+  dh::XGBDeviceAllocator<char> alloc;
+  thrust::stable_sort_by_key(thrust::cuda::par(alloc), keys.begin(), keys.end(), sorted_idx.begin(),
+                             [=] XGBOOST_DEVICE(Tup const &l, Tup const &r) {
+                               if (thrust::get<0>(l) != thrust::get<0>(r)) {
+                                 return thrust::get<0>(l) < thrust::get<0>(r);  // segment index
+                               }
+                               return thrust::get<1>(l) < thrust::get<1>(r);  // residue
+                             });
+}
+
 class CUDAStreamView;
 
 class CUDAEvent {
diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index fe91c8dce87e..9af9faf7c5e0 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -18,36 +18,6 @@
 
 namespace xgboost {
 namespace common {
-namespace detail {
-template <typename SegIt, typename ValIt>
-inline void SegmentedArgSort(SegIt seg_begin, SegIt seg_end, ValIt val_begin, ValIt val_end,
-                             dh::device_vector<size_t>* p_sorted_idx) {
-  using Tup = thrust::tuple<size_t, float>;
-  auto& sorted_idx = *p_sorted_idx;
-  size_t n = std::distance(val_begin, val_end);
-  sorted_idx.resize(n);
-  dh::Iota(dh::ToSpan(sorted_idx));
-  dh::device_vector<Tup> keys(sorted_idx.size());
-  auto key_it = dh::MakeTransformIterator<Tup>(
-      thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> Tup {
-        auto leaf_idx = dh::SegmentId(seg_begin, seg_end, i);
-        auto residue = val_begin[i];
-        return thrust::make_tuple(leaf_idx, residue);
-      });
-  dh::XGBCachingDeviceAllocator<char> caching;
-  thrust::copy(thrust::cuda::par(caching), key_it, key_it + keys.size(), keys.begin());
-
-  dh::XGBDeviceAllocator<char> alloc;
-  thrust::stable_sort_by_key(thrust::cuda::par(alloc), keys.begin(), keys.end(), sorted_idx.begin(),
-                             [=] XGBOOST_DEVICE(Tup const& l, Tup const& r) {
-                               if (thrust::get<0>(l) != thrust::get<0>(r)) {
-                                 return thrust::get<0>(l) < thrust::get<0>(r);  // segment index
-                               }
-                               return thrust::get<1>(l) < thrust::get<1>(r);  // residue
-                             });
-}
-}  // namespace detail
-
 /**
  * \brief Compute segmented quantile on GPU.
  *
@@ -65,7 +35,7 @@ void SegmentedQuantile(Context const* ctx, double alpha, SegIt seg_begin, SegIt
 
   dh::device_vector<size_t> sorted_idx;
   using Tup = thrust::tuple<size_t, float>;
-  detail::SegmentedArgSort(seg_begin, seg_end, val_begin, val_end, &sorted_idx);
+  dh::SegmentedArgSort(seg_begin, seg_end, val_begin, val_end, &sorted_idx);
   auto n_segments = std::distance(seg_begin, seg_end) - 1;
   if (n_segments <= 0) {
     return;
@@ -100,6 +70,7 @@ void SegmentedQuantile(Context const* ctx, double alpha, SegIt seg_begin, SegIt
     double x = alpha * static_cast<double>(n + 1);
     double k = std::floor(x) - 1;
     double d = (x - 1) - k;
+
     auto v0 = val[begin + static_cast<size_t>(k)];
     auto v1 = val[begin + static_cast<size_t>(k) + 1];
     d_results[seg_idx] = v0 + d * (v1 - v0);
@@ -112,7 +83,7 @@ void SegmentedWeightedQuantile(Context const* ctx, double alpha, SegIt seg_beg,
                                HostDeviceVector<float>* quantiles) {
   CHECK(alpha >= 0 && alpha <= 1);
   dh::device_vector<size_t> sorted_idx;
-  detail::SegmentedArgSort(seg_beg, seg_end, val_begin, val_end, &sorted_idx);
+  dh::SegmentedArgSort(seg_beg, seg_end, val_begin, val_end, &sorted_idx);
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
   size_t n_samples = std::distance(w_begin, w_end);
   dh::device_vector<float> weights_cdf(n_samples);
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index cf0d56986a1e..478502e8fe4e 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -675,21 +675,25 @@ void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_
                       RegTree* p_tree) {
   auto& tree = *p_tree;
   auto& quantiles = *p_quantiles;
+  auto const& h_node_idx = row_index.node_idx.HostVector();
 
-  size_t n_leaf{row_index.node_idx.Size()};
+  size_t n_leaf{h_node_idx.size()};
   rabit::Allreduce<rabit::op::Max>(&n_leaf, 1);
   CHECK(quantiles.empty() || quantiles.size() == n_leaf);
   if (quantiles.empty()) {
     quantiles.resize(n_leaf);
   }
+  for (size_t i = 0; i < n_leaf; ++i) {
+    if (std::isnan(quantiles[i])) {
+      quantiles[i] = tree[h_node_idx[i]].LeafValue();
+    }
+  }
   // use the mean value
   rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
   auto world = rabit::GetWorldSize();
   std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
                  [&](float q) { return q / static_cast<double>(world); });
 
-  // fixme: verify this is correct for external memory
-  auto const& h_node_idx = row_index.node_idx.HostVector();
   for (size_t i = 0; i < row_index.node_idx.Size(); ++i) {
     auto nidx = h_node_idx[i];
     auto q = quantiles[i];
@@ -777,15 +781,12 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
         q = common::WeightedQuantile(alpha, iter, iter + h_row_set.size(), w_it);
       }
       if (std::isnan(q)) {
-        // Edge case in distributed training where in a local worker a leaf can have 0
-        // samples.
         CHECK(h_row_set.empty());
-        q = 0;
+        q = tree[nidx].LeafValue();
       }
       results.at(k) = q;
     });
 
-    // fixme: verify this is correct for external memory
     for (size_t i = 0; i < results.size(); ++i) {
       quantiles[i] += results[i];
     }
diff --git a/tests/cpp/objective/test_regression_obj.cc b/tests/cpp/objective/test_regression_obj.cc
index f561c939345d..72a4b7c8db78 100644
--- a/tests/cpp/objective/test_regression_obj.cc
+++ b/tests/cpp/objective/test_regression_obj.cc
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2017-2021 XGBoost contributors
+ * Copyright 2017-2022 XGBoost contributors
  */
 #include <gtest/gtest.h>
 #include <xgboost/objective.h>
@@ -423,4 +423,45 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteError)) {
   ASSERT_EQ(tree[1].LeafValue(), -1);
   ASSERT_EQ(tree[2].LeafValue(), -4);
 }
+
+TEST(Objective, DeclareUnifiedTest(AbsoluteErrorLeaf)) {
+  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:absoluteerror", &ctx)};
+  obj->Configure({});
+
+  MetaInfo info;
+  info.labels.Reshape(16, 1);
+  info.num_row_ = info.labels.Size();
+  CHECK_EQ(info.num_row_, 16);
+  auto h_labels = info.labels.HostView().Values();
+  std::iota(h_labels.begin(), h_labels.end(), 0);
+  HostDeviceVector<float> predt(h_labels.size());
+  auto& h_predt = predt.HostVector();
+  for (size_t i = 0; i < h_predt.size(); ++i) {
+    h_predt[i] = h_labels[i] + i;
+  }
+
+  std::vector<RowIndexCache> row_idx_v;
+  row_idx_v.emplace_back(&ctx, info.labels.Shape(0));
+
+  auto& part = row_idx_v.back();
+  part.node_idx = {3, 4, 5, 6};
+  // starting from 3 to emulate subsampling, empty leaaft for node 4.
+  part.node_ptr = {3, 8, 8, 13, 16};
+  auto& h_row_idx = part.row_index.HostVector();
+  std::iota(h_row_idx.begin(), h_row_idx.end(), 0);
+
+  RegTree tree;
+  tree.ExpandNode(0, /*split_index=*/1, 2, true, 0.0f, 2.f, 3.f, 4.f, 2.f, 1.f, 1.f);
+  tree.ExpandNode(1, /*split_index=*/1, 2, true, 0.0f, 2.f, 3.f, 4.f, 2.f, 1.f, 1.f);
+  tree.ExpandNode(2, /*split_index=*/1, 2, true, 0.0f, 2.f, 3.f, 4.f, 2.f, 1.f, 1.f);
+  ASSERT_EQ(tree.GetNumLeaves(), 4);
+
+  auto empty_leaf = tree[4].LeafValue();
+  obj->UpdateTreeLeaf(row_idx_v, info, predt, 0, &tree);
+  ASSERT_EQ(tree[3].LeafValue(), -5);
+  ASSERT_EQ(tree[4].LeafValue(), empty_leaf);
+  ASSERT_EQ(tree[5].LeafValue(), -10);
+  ASSERT_EQ(tree[6].LeafValue(), -14);
+}
 }  // namespace xgboost

From 598c0e4473ee39af623b2f25600dcda9be282212 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 23:03:30 +0800
Subject: [PATCH 083/124] rename.

---
 src/common/stats.cuh | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index 9af9faf7c5e0..9983e3d72151 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -85,8 +85,8 @@ void SegmentedWeightedQuantile(Context const* ctx, double alpha, SegIt seg_beg,
   dh::device_vector<size_t> sorted_idx;
   dh::SegmentedArgSort(seg_beg, seg_end, val_begin, val_end, &sorted_idx);
   auto d_sorted_idx = dh::ToSpan(sorted_idx);
-  size_t n_samples = std::distance(w_begin, w_end);
-  dh::device_vector<float> weights_cdf(n_samples);
+  size_t n_weights = std::distance(w_begin, w_end);
+  dh::device_vector<float> weights_cdf(n_weights);
 
   dh::XGBCachingDeviceAllocator<char> caching;
   auto scan_key = dh::MakeTransformIterator<size_t>(
@@ -95,7 +95,7 @@ void SegmentedWeightedQuantile(Context const* ctx, double alpha, SegIt seg_beg,
   auto scan_val = dh::MakeTransformIterator<float>(
       thrust::make_counting_iterator(0ul),
       [=] XGBOOST_DEVICE(size_t i) { return w_begin[d_sorted_idx[i]]; });
-  thrust::inclusive_scan_by_key(thrust::cuda::par(caching), scan_key, scan_key + n_samples,
+  thrust::inclusive_scan_by_key(thrust::cuda::par(caching), scan_key, scan_key + n_weights,
                                 scan_val, weights_cdf.begin());
 
   auto n_segments = std::distance(seg_beg, seg_end) - 1;

From 19934f3d47b99e74b69dfd575bd8ddefde33e1db Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 23:09:39 +0800
Subject: [PATCH 084/124] Remove target.

---
 include/xgboost/objective.h                |  4 +---
 src/gbm/gbtree.cc                          |  4 +---
 src/objective/regression_obj.cu            | 19 +++++++++----------
 tests/cpp/objective/test_regression_obj.cc |  4 ++--
 4 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index a0dcd25c9795..91d33eaaf3ce 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -97,12 +97,10 @@ class ObjFunction : public Configurable {
    *
    * \param row_index The index of rows for each output leaf.
    * \param info MetaInfo providing labels and weights.
-   * \param target The index for target if we are training multi-target models, 0 otherwise.
    * \param p_tree Tree that needs to be updated.
    */
   virtual void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
-                              HostDeviceVector<float> const& prediction, uint32_t target,
-                              RegTree* p_tree) const {}
+                              HostDeviceVector<float> const& prediction, RegTree* p_tree) const {}
 
   /*!
    * \brief Create an objective function according to name.
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index 2a0371a3ce09..e71508ab5fc7 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -224,12 +224,10 @@ void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const
     return;
   }
   auto& trees = *p_trees;
-  auto targets = obj->Targets(p_fmat->Info());
   for (size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
     auto row_idx = updaters_.back()->GetRowIndexCache(tree_idx);
     // distinguish the difference between multi-class and multi-target.
-    auto target = targets > 1 ? gidx : 0;
-    obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), predictions, target, trees[tree_idx].get());
+    obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), predictions, trees[tree_idx].get());
   }
 }
 
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 478502e8fe4e..87de744364f5 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -669,7 +669,6 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .describe("Tweedie regression for insurance data.")
 .set_body([]() { return new TweedieRegression(); });
 
-
 namespace detail {
 void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_index,
                       RegTree* p_tree) {
@@ -704,8 +703,8 @@ void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_
 
 #if defined(XGBOOST_USE_CUDA)
 void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const> row_index,
-                          MetaInfo const& info, HostDeviceVector<float> const& predt,
-                          uint32_t target, float alpha, RegTree* p_tree) {
+                          MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
+                          RegTree* p_tree) {
   dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
   CHECK_EQ(row_index.size(), 1)
       << "External memory with GPU hist should have only 1 row partition.";
@@ -745,8 +744,8 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
 #endif  // defined(XGBOOST_USE_CUDA)
 
 void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
-                        MetaInfo const& info, HostDeviceVector<float> const& predt, uint32_t target,
-                        float alpha, RegTree* p_tree) {
+                        MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
+                        RegTree* p_tree) {
   auto& tree = *p_tree;
   CHECK(!row_index.empty());
   std::vector<float> quantiles(row_index.front().node_idx.Size(), 0);
@@ -761,7 +760,8 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
       CHECK_LT(k + 1, h_node_ptr.size());
       size_t n = h_node_ptr[k + 1] - h_node_ptr[k];
       auto h_row_set = part.row_index.HostSpan().subspan(h_node_ptr[k], n);
-      auto h_labels = info.labels.HostView().Slice(linalg::All(), target);
+      // multi-target not yet supported.
+      auto h_labels = info.labels.HostView().Slice(linalg::All(), 0);
       auto const& h_predt = predt.ConstHostVector();
       auto h_weights = linalg::MakeVec(&info.weights_);
 
@@ -830,13 +830,12 @@ class MeanAbsoluteError : public ObjFunction {
   }
 
   void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
-                      HostDeviceVector<float> const& prediction, uint32_t target,
-                      RegTree* p_tree) const override {
+                      HostDeviceVector<float> const& prediction, RegTree* p_tree) const override {
     if (ctx_->IsCPU()) {
-      detail::UpdateTreeLeafHost(ctx_, row_index, info, prediction, target, 0.5, p_tree);
+      detail::UpdateTreeLeafHost(ctx_, row_index, info, prediction, 0.5, p_tree);
     } else {
 #if defined(XGBOOST_USE_CUDA)
-      detail::UpdateTreeLeafDevice(ctx_, row_index, info, prediction, target, 0.5, p_tree);
+      detail::UpdateTreeLeafDevice(ctx_, row_index, info, prediction, 0.5, p_tree);
 #else
       common::AssertGPUSupport();
 #endif  //  defined(XGBOOST_USE_CUDA)
diff --git a/tests/cpp/objective/test_regression_obj.cc b/tests/cpp/objective/test_regression_obj.cc
index 72a4b7c8db78..96242b8d798e 100644
--- a/tests/cpp/objective/test_regression_obj.cc
+++ b/tests/cpp/objective/test_regression_obj.cc
@@ -419,7 +419,7 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteError)) {
   for (size_t i = 0; i < h_predt.size(); ++i) {
     h_predt[i] = labels[i] + i;
   }
-  obj->UpdateTreeLeaf(common::Span<RowIndexCache const>{row_idx}, info, predt, 0, &tree);
+  obj->UpdateTreeLeaf(common::Span<RowIndexCache const>{row_idx}, info, predt, &tree);
   ASSERT_EQ(tree[1].LeafValue(), -1);
   ASSERT_EQ(tree[2].LeafValue(), -4);
 }
@@ -458,7 +458,7 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteErrorLeaf)) {
   ASSERT_EQ(tree.GetNumLeaves(), 4);
 
   auto empty_leaf = tree[4].LeafValue();
-  obj->UpdateTreeLeaf(row_idx_v, info, predt, 0, &tree);
+  obj->UpdateTreeLeaf(row_idx_v, info, predt, &tree);
   ASSERT_EQ(tree[3].LeafValue(), -5);
   ASSERT_EQ(tree[4].LeafValue(), empty_leaf);
   ASSERT_EQ(tree[5].LeafValue(), -10);

From 9dff694f540ac9ef1c028c042f530c694edb5669 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 23:35:36 +0800
Subject: [PATCH 085/124] note.

---
 include/xgboost/objective.h            |  4 ++++
 tests/python-gpu/test_gpu_with_dask.py | 13 ++++++-------
 tests/python/test_with_dask.py         | 16 +++++++---------
 3 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 91d33eaaf3ce..40590402ae5c 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -95,6 +95,10 @@ class ObjFunction : public Configurable {
    * \brief Update the leaf values after a tree is built. Needed for objectives with 0
    *        hessian.
    *
+   *   Note that the leaf update is not well defined for distributed training as XGBoost
+   *   computes only an average of quantile between workers. This breaks when some leaf
+   *   have no sample assigned in a local worker.
+   *
    * \param row_index The index of rows for each output leaf.
    * \param info MetaInfo providing labels and weights.
    * \param p_tree Tree that needs to be updated.
diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
index 2ea4157c9045..30d1bad8f71b 100644
--- a/tests/python-gpu/test_gpu_with_dask.py
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -1,7 +1,7 @@
 """Copyright 2019-2022 XGBoost contributors"""
 import sys
 import os
-from typing import Type, TypeVar, Any, Dict, List, Tuple
+from typing import Type, TypeVar, Any, Dict, List
 import pytest
 import numpy as np
 import asyncio
@@ -9,7 +9,7 @@
 import subprocess
 from collections import OrderedDict
 from inspect import signature
-from hypothesis import given, strategies, settings, note, reproduce_failure
+from hypothesis import given, strategies, settings, note
 from hypothesis._settings import duration
 from test_gpu_updaters import parameter_strategy
 
@@ -201,9 +201,10 @@ def run_gpu_hist(
     )["history"]
     note(history)
 
-    # the approximate quantile for leaf can cause error on distributed training
-    if dataset.objective.endswith("-l1"):
-        assert tm.non_increasing(history["train"][dataset.metric], tolerance=1e-3)
+    # See note on `ObjFunction::UpdateTreeLeaf.
+    update_leaf = dataset.objective.endswith("-l1")
+    if update_leaf:
+        assert history[0] > history[-1]
     else:
         assert tm.non_increasing(history["train"][dataset.metric])
 
@@ -241,8 +242,6 @@ def test_dask_dataframe(self, local_cuda_cluster: LocalCUDACluster) -> None:
     @pytest.mark.parametrize(
         "local_cuda_cluster", [{"n_workers": 2}], indirect=["local_cuda_cluster"]
     )
-    @reproduce_failure('6.14.8', b'AXicY2ZgYGBiwAGYGBkYAADJAAk=')
-    # @reproduce_failure('6.14.8', b'AXicY2ZgYGBmZMAOGIESAAD+AAo=')
     def test_gpu_hist(
         self,
         params: Dict,
diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index 0e4cc1405615..cd5eb0570429 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -18,7 +18,7 @@
 import os
 import subprocess
 import hypothesis
-from hypothesis import given, settings, note, HealthCheck, reproduce_failure
+from hypothesis import given, settings, note, HealthCheck
 from test_updaters import hist_parameter_strategy, exact_parameter_strategy
 from test_with_sklearn import run_feature_weights, run_data_initialization
 from test_predict import verify_leaf_output
@@ -1289,10 +1289,11 @@ def is_stump():
         def minimum_bin():
             return "max_bin" in params and params["max_bin"] == 2
 
-        if minimum_bin() and is_stump():
-            assert tm.non_increasing(history, tolerance=1e-3)
-        elif dataset.objective.endswith("-l1"):
-            # the approximate quantile for leaf can cause error on distributed training
+        # See note on `ObjFunction::UpdateTreeLeaf.
+        update_leaf = dataset.objective.endswith("-l1")
+        if update_leaf:
+            assert history[0] > history[-1]
+        elif minimum_bin() and is_stump():
             assert tm.non_increasing(history, tolerance=1e-3)
         else:
             assert tm.non_increasing(history)
@@ -1302,7 +1303,6 @@ def minimum_bin():
     @given(params=hist_parameter_strategy,
            dataset=tm.dataset_strategy)
     @settings(deadline=None, suppress_health_check=suppress, print_blob=True)
-    @reproduce_failure('6.27.2', b'AAAAAAAAAAEAAA==')
     def test_hist(
             self, params: Dict, dataset: tm.TestDataset, client: "Client"
     ) -> None:
@@ -1312,10 +1312,8 @@ def test_hist(
     @given(params=exact_parameter_strategy,
            dataset=tm.dataset_strategy)
     @settings(deadline=None, suppress_health_check=suppress, print_blob=True)
-    # @reproduce_failure('6.27.2', b'AXicY2JnNDpff6rLhjXq6iuBo0IcVjvut1lMgdlLZwVcYuBofTkvQOMDKgAyAcAGTcPCw==')
-    # @reproduce_failure('6.36.1', b'AXicY2BkIBUwArUAAAB0AAQ=')
     def test_approx(
-            self, client: "Client", params: Dict, dataset: tm.TestDataset
+        self, client: "Client", params: Dict, dataset: tm.TestDataset
     ) -> None:
         num_rounds = 30
         # params["eta"] = 0.1

From 6bbba67fd5edcb02afa506ae8e44388704ffa6ff Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 15 Apr 2022 23:39:51 +0800
Subject: [PATCH 086/124] test.

---
 tests/python-gpu/test_gpu_with_dask.py | 5 +++--
 tests/python/test_with_dask.py         | 5 +++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
index 30d1bad8f71b..82d4a009d88e 100644
--- a/tests/python-gpu/test_gpu_with_dask.py
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -202,9 +202,10 @@ def run_gpu_hist(
     note(history)
 
     # See note on `ObjFunction::UpdateTreeLeaf.
-    update_leaf = dataset.objective.endswith("-l1")
+    update_leaf = dataset.name.endswith("-l1")
     if update_leaf:
-        assert history[0] > history[-1]
+        assert history[0] >= history[-1]
+        return
     else:
         assert tm.non_increasing(history["train"][dataset.metric])
 
diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index cd5eb0570429..6a63488034b4 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -1290,9 +1290,10 @@ def minimum_bin():
             return "max_bin" in params and params["max_bin"] == 2
 
         # See note on `ObjFunction::UpdateTreeLeaf.
-        update_leaf = dataset.objective.endswith("-l1")
+        update_leaf = dataset.name.endswith("-l1")
         if update_leaf:
-            assert history[0] > history[-1]
+            assert history[0] >= history[-1]
+            return
         elif minimum_bin() and is_stump():
             assert tm.non_increasing(history, tolerance=1e-3)
         else:

From e58a3394bf4dc3547db2406784943b49a43805f3 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 00:07:10 +0800
Subject: [PATCH 087/124] CPU build.

---
 src/objective/regression_obj.cu | 1 +
 src/tree/updater_gpu_hist.cu    | 5 +++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 87de744364f5..fc116fcb08aa 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -6,6 +6,7 @@
  */
 
 #include <dmlc/omp.h>
+#include <rabit/rabit.h>
 #include <xgboost/logging.h>
 #include <xgboost/objective.h>
 #include <xgboost/tree_model.h>
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 4b5f00b57904..6c287fb85eab 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -223,8 +223,9 @@ struct GPUHistMakerDevice {
     // Init histogram
     hist.Init(ctx_->gpu_id, page->Cuts().TotalBins());
     monitor.Init(std::string("GPUHistMakerDevice") + std::to_string(ctx_->gpu_id));
-    feature_groups.reset(new FeatureGroups(
-        page->Cuts(), page->is_dense, dh::MaxSharedMemoryOptin(ctx_->gpu_id), sizeof(GradientSumT)));
+    feature_groups.reset(new FeatureGroups(page->Cuts(), page->is_dense,
+                                           dh::MaxSharedMemoryOptin(ctx_->gpu_id),
+                                           sizeof(GradientSumT)));
   }
 
   ~GPUHistMakerDevice() {  // NOLINT

From 3faab6c5f8530fa22010734006e9ad4168c94526 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 00:47:47 +0800
Subject: [PATCH 088/124] Json schema.

---
 doc/model.schema | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/doc/model.schema b/doc/model.schema
index b192cabc6864..02725cb36d31 100644
--- a/doc/model.schema
+++ b/doc/model.schema
@@ -400,7 +400,6 @@
                 "reg_loss_param"
               ]
             },
-
             {
               "type": "object",
               "properties": {
@@ -433,6 +432,14 @@
                 "tweedie_regression_param"
               ]
             },
+            {
+              "properties": {
+                "name": {
+                  "const": "reg:absoluteerror"
+                }
+              },
+              "type": "object"
+            },
             {
               "type": "object",
               "properties": {

From 35b64dc8795c1b95923fba3f7db752077f0fede0 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 00:51:58 +0800
Subject: [PATCH 089/124] Fix tests.

---
 tests/python-gpu/test_gpu_prediction.py | 2 ++
 tests/python-gpu/test_gpu_with_dask.py  | 2 +-
 tests/python/test_with_dask.py          | 2 +-
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/python-gpu/test_gpu_prediction.py b/tests/python-gpu/test_gpu_prediction.py
index 38f4db07d366..29e6797e71f3 100644
--- a/tests/python-gpu/test_gpu_prediction.py
+++ b/tests/python-gpu/test_gpu_prediction.py
@@ -249,6 +249,8 @@ def predict_df(x):
            tm.dataset_strategy, shap_parameter_strategy)
     @settings(deadline=None, print_blob=True)
     def test_shap(self, num_rounds, dataset, param):
+        if dataset.name.endswith("-l1"):  # not supported by exact tree method
+            return
         param.update({"predictor": "gpu_predictor", "gpu_id": 0})
         param = dataset.set_params(param)
         dmat = dataset.get_dmat()
diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
index 82d4a009d88e..10226dbed5ce 100644
--- a/tests/python-gpu/test_gpu_with_dask.py
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -203,7 +203,7 @@ def run_gpu_hist(
 
     # See note on `ObjFunction::UpdateTreeLeaf.
     update_leaf = dataset.name.endswith("-l1")
-    if update_leaf:
+    if update_leaf and len(history) >= 2:
         assert history[0] >= history[-1]
         return
     else:
diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index 6a63488034b4..8566edb172b8 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -1291,7 +1291,7 @@ def minimum_bin():
 
         # See note on `ObjFunction::UpdateTreeLeaf.
         update_leaf = dataset.name.endswith("-l1")
-        if update_leaf:
+        if update_leaf and len(history) >= 2:
             assert history[0] >= history[-1]
             return
         elif minimum_bin() and is_stump():

From 1105b6fb2818fe8bd604aa767ac4aeb01da05156 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 01:26:23 +0800
Subject: [PATCH 090/124] Skip tests.

---
 tests/python-gpu/test_gpu_prediction.py | 4 +++-
 tests/python-gpu/test_gpu_with_dask.py  | 4 ++--
 tests/python/test_updaters.py           | 2 ++
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/tests/python-gpu/test_gpu_prediction.py b/tests/python-gpu/test_gpu_prediction.py
index 29e6797e71f3..4e41e637f7de 100644
--- a/tests/python-gpu/test_gpu_prediction.py
+++ b/tests/python-gpu/test_gpu_prediction.py
@@ -249,7 +249,7 @@ def predict_df(x):
            tm.dataset_strategy, shap_parameter_strategy)
     @settings(deadline=None, print_blob=True)
     def test_shap(self, num_rounds, dataset, param):
-        if dataset.name.endswith("-l1"):  # not supported by exact tree method
+        if dataset.name.endswith("-l1"):  # not supported by the exact tree method
             return
         param.update({"predictor": "gpu_predictor", "gpu_id": 0})
         param = dataset.set_params(param)
@@ -265,6 +265,8 @@ def test_shap(self, num_rounds, dataset, param):
            tm.dataset_strategy, shap_parameter_strategy)
     @settings(deadline=None, max_examples=20, print_blob=True)
     def test_shap_interactions(self, num_rounds, dataset, param):
+        if dataset.name.endswith("-l1"):  # not supported by the exact tree method
+            return
         param.update({"predictor": "gpu_predictor", "gpu_id": 0})
         param = dataset.set_params(param)
         dmat = dataset.get_dmat()
diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
index 10226dbed5ce..105b15c46813 100644
--- a/tests/python-gpu/test_gpu_with_dask.py
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -198,7 +198,7 @@ def run_gpu_hist(
         dtrain=m,
         num_boost_round=num_rounds,
         evals=[(m, "train")],
-    )["history"]
+    )["history"]["train"]
     note(history)
 
     # See note on `ObjFunction::UpdateTreeLeaf.
@@ -207,7 +207,7 @@ def run_gpu_hist(
         assert history[0] >= history[-1]
         return
     else:
-        assert tm.non_increasing(history["train"][dataset.metric])
+        assert tm.non_increasing(history[dataset.metric])
 
 
 @pytest.mark.skipif(**tm.no_cudf())
diff --git a/tests/python/test_updaters.py b/tests/python/test_updaters.py
index cdf40d843b1a..4b56d37d4493 100644
--- a/tests/python/test_updaters.py
+++ b/tests/python/test_updaters.py
@@ -40,6 +40,8 @@ class TestTreeMethod:
            tm.dataset_strategy)
     @settings(deadline=None, print_blob=True)
     def test_exact(self, param, num_rounds, dataset):
+        if dataset.name.endswith("-l1"):
+            return
         param['tree_method'] = 'exact'
         param = dataset.set_params(param)
         result = train_result(param, dataset.get_dmat(), num_rounds)

From 3a25702337f82d53cf6c26b53b3022c1bd5c69fd Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 02:04:54 +0800
Subject: [PATCH 091/124] Skip test.

---
 tests/python-gpu/test_gpu_with_dask.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
index 105b15c46813..f4db033c81f0 100644
--- a/tests/python-gpu/test_gpu_with_dask.py
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -198,7 +198,7 @@ def run_gpu_hist(
         dtrain=m,
         num_boost_round=num_rounds,
         evals=[(m, "train")],
-    )["history"]["train"]
+    )["history"]["train"][dataset.metric]
     note(history)
 
     # See note on `ObjFunction::UpdateTreeLeaf.
@@ -207,7 +207,7 @@ def run_gpu_hist(
         assert history[0] >= history[-1]
         return
     else:
-        assert tm.non_increasing(history[dataset.metric])
+        assert tm.non_increasing(history)
 
 
 @pytest.mark.skipif(**tm.no_cudf())

From 4a7359c9f2a47e7d016e77ed01674a32f2b35fa4 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 16:06:09 +0800
Subject: [PATCH 092/124] Cleanup.

---
 include/xgboost/tree_model.h          |  8 ++++++++
 src/common/common.h                   | 15 ---------------
 src/common/stats.cuh                  |  1 -
 src/gbm/gbtree.cc                     |  1 -
 tests/cpp/common/test_stats.cc        |  2 +-
 tests/python-gpu/test_gpu_updaters.py |  2 +-
 6 files changed, 10 insertions(+), 19 deletions(-)

diff --git a/include/xgboost/tree_model.h b/include/xgboost/tree_model.h
index 789d336f59fb..7912da809378 100644
--- a/include/xgboost/tree_model.h
+++ b/include/xgboost/tree_model.h
@@ -736,11 +736,19 @@ inline bool RegTree::FVec::HasMissing() const {
   return has_missing_;
 }
 
+/**
+ * \brief A cache for row partition, each partition is the row index of a tree leaf.
+ */
 struct RowIndexCache {
   HostDeviceVector<size_t> row_index;
   HostDeviceVector<size_t> node_ptr;
   HostDeviceVector<bst_node_t> node_idx;
 
+  /**
+   * \param ctx Context
+   * \param n_samples The number of samples for this cache, which equals to the number of
+   *                  samples in a single page from DMatarix.
+   */
   RowIndexCache(Context const* ctx, size_t n_samples) {
     if (!ctx->IsCPU()) {
       row_index.SetDevice(ctx->gpu_id);
diff --git a/src/common/common.h b/src/common/common.h
index 95b7e4103261..d909a37691ae 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -251,15 +251,6 @@ std::vector<Idx> ArgSort(Container const &array, Comp comp = std::less<V>{}) {
   return result;
 }
 
-template <typename V, typename Comp = std::less<V>>
-std::vector<size_t> ArgSort(linalg::TensorView<V, 1> array, Comp comp = std::less<V>{}) {
-  std::vector<size_t> result(array.Size());
-  std::iota(result.begin(), result.end(), 0);
-  auto op = [&array, comp](size_t const &l, size_t const &r) { return comp(array(l), array(r)); };
-  XGBOOST_PARALLEL_STABLE_SORT(result.begin(), result.end(), op);
-  return result;
-}
-
 struct OptionalWeights {
   Span<float const> weights;
   float dft{1.0f};  // fixme: make this compile time constant
@@ -270,15 +261,9 @@ struct OptionalWeights {
   XGBOOST_DEVICE float operator[](size_t i) const { return weights.empty() ? dft : weights[i]; }
 };
 
-
 /**
  * Last index of a group in a CSR style of index pointer.
  */
-template <typename Idx>
-XGBOOST_DEVICE size_t LastOf(size_t group, common::Span<Idx> indptr) {
-  return indptr[group + 1] - 1;
-}
-
 template <typename Indexable>
 XGBOOST_DEVICE size_t LastOf(size_t group, Indexable const &indptr) {
   return indptr[group + 1] - 1;
diff --git a/src/common/stats.cuh b/src/common/stats.cuh
index 9983e3d72151..9d9e526a8576 100644
--- a/src/common/stats.cuh
+++ b/src/common/stats.cuh
@@ -6,7 +6,6 @@
 
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/permutation_iterator.h>
-#include <thrust/sort.h>
 
 #include <iterator>  // std::distance
 
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index e71508ab5fc7..2dade8720547 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -226,7 +226,6 @@ void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const
   auto& trees = *p_trees;
   for (size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
     auto row_idx = updaters_.back()->GetRowIndexCache(tree_idx);
-    // distinguish the difference between multi-class and multi-target.
     obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), predictions, trees[tree_idx].get());
   }
 }
diff --git a/tests/cpp/common/test_stats.cc b/tests/cpp/common/test_stats.cc
index 121dce4cba09..2a1e375c0f20 100644
--- a/tests/cpp/common/test_stats.cc
+++ b/tests/cpp/common/test_stats.cc
@@ -8,7 +8,7 @@
 
 namespace xgboost {
 namespace common {
-TEST(Stats, Percentil) {
+TEST(Stats, Quantile) {
   {
     linalg::Tensor<float, 1> arr({20.f, 0.f, 15.f, 50.f, 40.f, 0.f, 35.f}, {7}, Context::kCpuId);
     std::vector<size_t> index{0, 2, 3, 4, 6};
diff --git a/tests/python-gpu/test_gpu_updaters.py b/tests/python-gpu/test_gpu_updaters.py
index 9eb68a2d10a1..c6574bca4a22 100644
--- a/tests/python-gpu/test_gpu_updaters.py
+++ b/tests/python-gpu/test_gpu_updaters.py
@@ -3,7 +3,7 @@
 import gc
 import pytest
 import xgboost as xgb
-from hypothesis import given, strategies, assume, settings, note, reproduce_failure, seed
+from hypothesis import given, strategies, assume, settings, note
 
 sys.path.append("tests/python")
 import testing as tm

From b2b79f4752d204e20c1e3332f2ec045c8ee5f05f Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 17:40:38 +0800
Subject: [PATCH 093/124] Make optional.

---
 src/tree/updater_approx.cc        | 3 +++
 src/tree/updater_quantile_hist.cc | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index aa9de4f71858..a35518a07db7 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -158,6 +158,9 @@ class GloablApproxBuilder {
   void LeafPartition(RegTree const &tree, common::Span<float> hess,
                      std::vector<RowIndexCache> *p_out_row_indices) {
     monitor_->Start(__func__);
+    if (!evaluator_.Task().UpdateTreeLeaf()) {
+      return;
+    }
     CHECK(p_out_row_indices->empty());
     for (auto const &part : partitioner_) {
       part.LeafPartition(ctx_, tree, hess, p_out_row_indices);
diff --git a/src/tree/updater_quantile_hist.cc b/src/tree/updater_quantile_hist.cc
index 15aff5a16c01..b2f01e3b7917 100644
--- a/src/tree/updater_quantile_hist.cc
+++ b/src/tree/updater_quantile_hist.cc
@@ -177,6 +177,9 @@ void QuantileHistMaker::Builder<GradientSumT>::LeafPartition(
     RegTree const &tree, common::Span<GradientPair const> gpair,
     std::vector<RowIndexCache> *p_out_row_indices) {
   monitor_->Start(__func__);
+  if (!evaluator_->Task().UpdateTreeLeaf()) {
+    return;
+  }
   CHECK(p_out_row_indices->empty());
   for (auto const &part : partitioner_) {
     part.LeafPartition(ctx_, tree, gpair, p_out_row_indices);

From 5cc1a4a465b6524e586588c2a43a4b0a5e556377 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 18:09:27 +0800
Subject: [PATCH 094/124] Allreduce.

---
 doc/parameter.rst                     |  1 +
 src/objective/regression_obj.cu       | 44 ++++++++++++++++++++-------
 src/tree/gpu_hist/row_partitioner.cuh | 13 +++++---
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/doc/parameter.rst b/doc/parameter.rst
index 781150490082..b361b01d4d9f 100644
--- a/doc/parameter.rst
+++ b/doc/parameter.rst
@@ -349,6 +349,7 @@ Specify the learning task and the corresponding learning objective. The objectiv
   - ``reg:squaredlogerror``: regression with squared log loss :math:`\frac{1}{2}[log(pred + 1) - log(label + 1)]^2`.  All input labels are required to be greater than -1.  Also, see metric ``rmsle`` for possible issue  with this objective.
   - ``reg:logistic``: logistic regression.
   - ``reg:pseudohubererror``: regression with Pseudo Huber loss, a twice differentiable alternative to absolute loss.
+  - ``reg:absoluteerror``: Regression with L1 error. When tree model is used, leaf value is refreshed after tree construction.
   - ``binary:logistic``: logistic regression for binary classification, output probability
   - ``binary:logitraw``: logistic regression for binary classification, output score before logistic transformation
   - ``binary:hinge``: hinge loss for binary classification. This makes predictions of 0 or 1, rather than producing probabilities.
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index fc116fcb08aa..58cb60c235ea 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -683,22 +683,31 @@ void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_
   if (quantiles.empty()) {
     quantiles.resize(n_leaf);
   }
+
+  // number of workers that have valid quantiles
+  std::vector<int32_t> n_valids(quantiles.size());
+  std::transform(quantiles.cbegin(), quantiles.cend(), n_valids.begin(),
+                 [](float q) { return static_cast<int32_t>(!std::isnan(q)); });
+  rabit::Allreduce<rabit::op::Sum>(n_valids.data(), n_valids.size());
+  // convert to 0 for all reduce
+  std::replace_if(
+      quantiles.begin(), quantiles.end(), [](float q) { return std::isnan(q); }, 0.f);
+  // use the mean value
+  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
   for (size_t i = 0; i < n_leaf; ++i) {
-    if (std::isnan(quantiles[i])) {
+    if (n_valids[i] > 0) {
+      quantiles[i] /= static_cast<double>(n_valids[i]);
+    } else {
+      // Use original leaf value if no worker can provide the quantile.
       quantiles[i] = tree[h_node_idx[i]].LeafValue();
     }
   }
-  // use the mean value
-  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
-  auto world = rabit::GetWorldSize();
-  std::transform(quantiles.begin(), quantiles.end(), quantiles.begin(),
-                 [&](float q) { return q / static_cast<double>(world); });
 
   for (size_t i = 0; i < row_index.node_idx.Size(); ++i) {
     auto nidx = h_node_idx[i];
     auto q = quantiles[i];
     CHECK(tree[nidx].IsLeaf());
-    tree[nidx].SetLeaf(q);  // fixme: exact tree method
+    tree[nidx].SetLeaf(q);
   }
 }
 
@@ -750,11 +759,14 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
   auto& tree = *p_tree;
   CHECK(!row_index.empty());
   std::vector<float> quantiles(row_index.front().node_idx.Size(), 0);
+  std::vector<int32_t> n_valids(quantiles.size(), 0);
+  // loop over external memory partitions
   for (auto const& part : row_index) {
     std::vector<float> results(part.node_idx.Size());
     auto const& h_node_idx = part.node_idx.ConstHostVector();
     auto const& h_node_ptr = part.node_ptr.ConstHostVector();
     CHECK_LE(h_node_ptr.back(), info.num_row_);
+    // loop over each leaf
     common::ParallelFor(results.size(), ctx->Threads(), [&](size_t k) {
       auto nidx = h_node_idx[k];
       CHECK(tree[nidx].IsLeaf());
@@ -783,17 +795,27 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
       }
       if (std::isnan(q)) {
         CHECK(h_row_set.empty());
-        q = tree[nidx].LeafValue();
       }
       results.at(k) = q;
     });
 
+    // sum result from each external memory partition to quantiles
     for (size_t i = 0; i < results.size(); ++i) {
-      quantiles[i] += results[i];
+      if (!std::isnan(results[i])) {
+        quantiles[i] += results[i];
+        n_valids[i]++;
+      }
+    }
+  }
+
+  for (size_t i = 0; i < quantiles.size(); ++i) {
+    if (n_valids[i] > 0) {
+      quantiles[i] /= n_valids[i];
+    } else {
+      // mark that no page has valid sample in the i^th leaf
+      quantiles[i] = std::numeric_limits<float>::quiet_NaN();
     }
   }
-  std::transform(quantiles.cbegin(), quantiles.cend(), quantiles.begin(),
-                 [&](float q) { return q / static_cast<float>(row_index.size()); });
 
   UpdateLeafValues(&quantiles, row_index.front(), p_tree);
 }
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 5fa0318fd0bf..7fed5baff75d 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -175,13 +175,16 @@ class RowPartitioner {
   }
 
   /**
-   * \brief Finalise the position of all training instances after tree
-   * construction is complete. Does not update any other meta information in
-   * this data structure, so should only be used at the end of training.
+   * \brief Finalise the position of all training instances after tree construction is
+   * complete. Does not update any other meta information in this data structure, so
+   * should only be used at the end of training.
+   *
+   *   When the task requires update leaf, this function will copy the row partitions into
+   *   p_out_row_indices. Note that the node ptr might not start from 0 due to sampling.
    *
    * \param p_out_row_indices Row partitions for each leaf.
    * \param op Device lambda. Should provide the row index and current position as an
-   * argument and return the new position for this training instance.
+   *           argument and return the new position for this training instance.
    */
   template <typename FinalisePositionOpT, typename Sampledp>
   void FinalisePosition(Context const* ctx, RegTree const* p_tree, size_t n_leaf, ObjInfo task,
@@ -313,7 +316,7 @@ class RowPartitioner {
         return true;
       });
       CHECK_EQ(leaves.size(), n_leaf);
-      // Fill all the leaves that don't have any sample. This is hacky and inefficien. An
+      // Fill all the leaves that don't have any sample. This is hacky and inefficient. An
       // alternative is to leave the objective to handle missing leaf, which is more messy
       // as we need to take other distributed workers into account.
       detail::FillMissingLeaf(leaves, &row_indices);

From 493795b35fb5ce0ed183f195e605bf847e2d85be Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 19:07:01 +0800
Subject: [PATCH 095/124] Test.

---
 src/common/partition_builder.h                | 31 ++++-----
 tests/cpp/common/test_partition_builder.cc    |  4 +-
 .../cpp/tree/gpu_hist/test_row_partitioner.cu |  2 +-
 tests/cpp/tree/test_approx.cc                 | 67 ++++++++++++++++++-
 4 files changed, 86 insertions(+), 18 deletions(-)

diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index 44f6a0ba464b..0cf929a78c99 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -285,31 +285,32 @@ class PartitionBuilder {
   void LeafPartition(Context const* ctx, RegTree const& tree, RowSetCollection const& row_set,
                      std::vector<RowIndexCache>* p_out_row_indices, Sampledp sampledp) const {
     p_out_row_indices->emplace_back(ctx, row_set.Data()->size());
-    auto& h_row_index = p_out_row_indices->back().row_index.HostVector();
-    auto& h_node_ptr = p_out_row_indices->back().node_ptr.HostVector();
-    auto& h_node_nidx = p_out_row_indices->back().node_idx.HostVector();
-    CHECK(h_node_ptr.empty());
+    auto& h_ridx = p_out_row_indices->back().row_index.HostVector();
+    auto& h_nptr = p_out_row_indices->back().node_ptr.HostVector();
+    auto& h_nidx = p_out_row_indices->back().node_idx.HostVector();
+    CHECK(h_nptr.empty());
 
     auto p_begin = row_set.Data()->data();
     size_t offset{0};
-    h_node_ptr.push_back(offset);
+    h_nptr.push_back(offset);
     for (auto node : row_set) {
-      if (!node.begin) {
+      if (!tree[node.node_id].IsLeaf()) {
         continue;
       }
-      CHECK(node.begin && tree[node.node_id].IsLeaf());
-      size_t ptr_offset = node.end - p_begin;
-      CHECK_LE(ptr_offset, row_set.Data()->size()) << node.node_id;
-      for (auto idx = node.begin; idx != node.end; ++idx) {
-        if (!sampledp(*idx)) {
-          h_row_index[offset++] = *idx;
+      if (node.begin) {  // guard for empty node.
+        size_t ptr_offset = node.end - p_begin;
+        CHECK_LE(ptr_offset, row_set.Data()->size()) << node.node_id;
+        for (auto idx = node.begin; idx != node.end; ++idx) {
+          if (!sampledp(*idx)) {
+            h_ridx[offset++] = *idx;
+          }
         }
       }
-      h_node_ptr.push_back(offset);
-      h_node_nidx.push_back(node.node_id);
+      h_nptr.push_back(offset);
+      h_nidx.push_back(node.node_id);
     }
     CHECK_LE(offset, row_set.Data()->size());
-    h_row_index.resize(offset);
+    h_ridx.resize(offset);
   }
 
  protected:
diff --git a/tests/cpp/common/test_partition_builder.cc b/tests/cpp/common/test_partition_builder.cc
index 885b924e71c1..aeedf8e16b54 100644
--- a/tests/cpp/common/test_partition_builder.cc
+++ b/tests/cpp/common/test_partition_builder.cc
@@ -1,3 +1,6 @@
+/*!
+ * Copyright 2021-2022 by XGBoost Contributors
+ */
 #include <gtest/gtest.h>
 #include <vector>
 #include <string>
@@ -74,6 +77,5 @@ TEST(PartitionBuilder, BasicTest) {
     ASSERT_EQ(n_right, (kBlockSize - rows_for_left_node[nid]) * tasks[nid]);
   }
 }
-
 }  // namespace common
 }  // namespace xgboost
diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index 96f464a6126a..32e6d2cf6a1f 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2019-2021 by XGBoost Contributors
+ * Copyright 2019-2022 by XGBoost Contributors
  */
 #include <gtest/gtest.h>
 #include <algorithm>
diff --git a/tests/cpp/tree/test_approx.cc b/tests/cpp/tree/test_approx.cc
index a37c0973627e..3aa29ddc623d 100644
--- a/tests/cpp/tree/test_approx.cc
+++ b/tests/cpp/tree/test_approx.cc
@@ -44,9 +44,9 @@ TEST(Approx, Partitioner) {
       float split_value = page.cut.Values().at(ptr / 2);
       RegTree tree;
       GetSplit(&tree, split_value, &candidates);
-      auto left_nidx = tree[RegTree::kRoot].LeftChild();
       partitioner.UpdatePosition(&ctx, page, candidates, &tree);
 
+      auto left_nidx = tree[RegTree::kRoot].LeftChild();
       auto elem = partitioner[left_nidx];
       ASSERT_LT(elem.Size(), n_samples);
       ASSERT_GT(elem.Size(), 1);
@@ -54,6 +54,7 @@ TEST(Approx, Partitioner) {
         auto value = page.cut.Values().at(page.index[*it]);
         ASSERT_LE(value, split_value);
       }
+
       auto right_nidx = tree[RegTree::kRoot].RightChild();
       elem = partitioner[right_nidx];
       for (auto it = elem.begin; it != elem.end; ++it) {
@@ -63,5 +64,69 @@ TEST(Approx, Partitioner) {
     }
   }
 }
+namespace {
+void TestLeafPartition(size_t n_samples) {
+  size_t const n_features = 2, base_rowid = 0;
+  common::RowSetCollection row_set;
+  ApproxRowPartitioner partitioner{n_samples, base_rowid};
+
+  auto Xy = RandomDataGenerator{n_samples, n_features, 0}.GenerateDMatrix(true);
+  GenericParameter ctx;
+  std::vector<CPUExpandEntry> candidates{{0, 0, 0.4}};
+  RegTree tree;
+  std::vector<float> hess(n_samples, 0);
+  // emulate sampling
+  auto not_sampled = [](size_t i) {
+    size_t const kSampleFactor{3};
+    return i % kSampleFactor != 0;
+  };
+  size_t n{0};
+  for (size_t i = 0; i < hess.size(); ++i) {
+    if (not_sampled(i)) {
+      hess[i] = 1.0f;
+      ++n;
+    }
+  }
+
+  std::vector<size_t> h_nptr;
+  float split_value;
+  for (auto const& page : Xy->GetBatches<GHistIndexMatrix>({Context::kCpuId, 64})) {
+    bst_feature_t const split_ind = 0;
+    auto ptr = page.cut.Ptrs()[split_ind + 1];
+    split_value = page.cut.Values().at(ptr / 2);
+    GetSplit(&tree, split_value, &candidates);
+    partitioner.UpdatePosition(&ctx, page, candidates, &tree);
+    std::vector<RowIndexCache> cache;
+    partitioner.LeafPartition(&ctx, tree, hess, &cache);
+    auto const& row_idx = cache.front();
+    ASSERT_EQ(n, row_idx.row_index.Size());
+    h_nptr = row_idx.node_ptr.ConstHostVector();
+    ASSERT_EQ(h_nptr.size(), 3);
+    ASSERT_EQ(h_nptr[0], 0);
+    ASSERT_EQ(h_nptr[2], n);  // equal to sampled rows
+
+    ASSERT_EQ(row_idx.node_idx.Size(), 2);
+    ASSERT_EQ(row_idx.node_idx.HostVector()[0], 1);
+    ASSERT_EQ(row_idx.node_idx.HostVector()[1], 2);
+  }
+
+  for (auto const& page : Xy->GetBatches<SparsePage>()) {
+    auto batch = page.GetView();
+    size_t left{0};
+    for (size_t i = 0; i < batch.Size(); ++i) {
+      if (not_sampled(i) && batch[i].front().fvalue < split_value) {
+        left++;
+      }
+    }
+    ASSERT_EQ(left, h_nptr[1]);  // equal to number of sampled assigned to left
+  }
+}
+}  // anonymous namespace
+
+TEST(Approx, LeafPartition) {
+  for (auto n_samples : {0ul, 1ul, 128ul, 256ul}) {
+    TestLeafPartition(n_samples);
+  }
+}
 }  // namespace tree
 }  // namespace xgboost

From 980219f2e728c13d852823c0d788c08bbaa02226 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 19:13:05 +0800
Subject: [PATCH 096/124] Header.

---
 include/xgboost/objective.h               | 2 +-
 plugin/example/custom_obj.cc              | 2 +-
 src/common/device_helpers.cuh             | 2 +-
 src/common/row_set.h                      | 2 +-
 src/gbm/gbtree.h                          | 2 +-
 src/metric/auc.cu                         | 2 +-
 src/objective/rank_obj.cu                 | 2 +-
 src/objective/regression_loss.h           | 2 +-
 src/tree/gpu_hist/row_partitioner.cuh     | 2 +-
 src/tree/updater_approx.h                 | 2 +-
 tests/cpp/gbm/test_gbtree.cc              | 2 +-
 tests/cpp/predictor/test_cpu_predictor.cc | 2 +-
 tests/cpp/tree/test_gpu_hist.cu           | 2 +-
 13 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 40590402ae5c..1fd3515c1cca 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2014-2019 by Contributors
+ * Copyright 2014-2022 by Contributors
  * \file objective.h
  * \brief interface of objective function used by xgboost.
  * \author Tianqi Chen, Kailong Chen
diff --git a/plugin/example/custom_obj.cc b/plugin/example/custom_obj.cc
index 1b26fea410a6..e220e4497141 100644
--- a/plugin/example/custom_obj.cc
+++ b/plugin/example/custom_obj.cc
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2015-2019 by Contributors
+ * Copyright 2015-2022 by Contributors
  * \file custom_metric.cc
  * \brief This is an example to define plugin of xgboost.
  *  This plugin defines the additional metric function.
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
index b4da11b559d5..2286db60d8de 100644
--- a/src/common/device_helpers.cuh
+++ b/src/common/device_helpers.cuh
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2017-2021 XGBoost contributors
+ * Copyright 2017-2022 XGBoost contributors
  */
 #pragma once
 #include <thrust/device_ptr.h>
diff --git a/src/common/row_set.h b/src/common/row_set.h
index 64f6089e8ec3..87d5f52874f2 100644
--- a/src/common/row_set.h
+++ b/src/common/row_set.h
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2017 by Contributors
+ * Copyright 2017-2022 by Contributors
  * \file row_set.h
  * \brief Quick Utility to compute subset of rows
  * \author Philip Cho, Tianqi Chen
diff --git a/src/gbm/gbtree.h b/src/gbm/gbtree.h
index 0e5343cec5c7..4c4b7eb93756 100644
--- a/src/gbm/gbtree.h
+++ b/src/gbm/gbtree.h
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2014-2021 by Contributors
+ * Copyright 2014-2022 by Contributors
  * \file gbtree.cc
  * \brief gradient boosted tree implementation.
  * \author Tianqi Chen
diff --git a/src/metric/auc.cu b/src/metric/auc.cu
index 52504ecbfbdf..5faa116c8561 100644
--- a/src/metric/auc.cu
+++ b/src/metric/auc.cu
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2021 by XGBoost Contributors
+ * Copyright 2021-2022 by XGBoost Contributors
  */
 #include <thrust/scan.h>
 #include <cub/cub.cuh>
diff --git a/src/objective/rank_obj.cu b/src/objective/rank_obj.cu
index 899529824320..f1c8702102df 100644
--- a/src/objective/rank_obj.cu
+++ b/src/objective/rank_obj.cu
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2015-2019 XGBoost contributors
+ * Copyright 2015-2022 XGBoost contributors
  */
 #include <dmlc/omp.h>
 #include <dmlc/timer.h>
diff --git a/src/objective/regression_loss.h b/src/objective/regression_loss.h
index cbc161e0840d..f394432a8f28 100644
--- a/src/objective/regression_loss.h
+++ b/src/objective/regression_loss.h
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2017-2019 XGBoost contributors
+ * Copyright 2017-2022 XGBoost contributors
  */
 #ifndef XGBOOST_OBJECTIVE_REGRESSION_LOSS_H_
 #define XGBOOST_OBJECTIVE_REGRESSION_LOSS_H_
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 7fed5baff75d..2ffb40322467 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2017-2019 XGBoost contributors
+ * Copyright 2017-2022 XGBoost contributors
  */
 #pragma once
 #include <limits>
diff --git a/src/tree/updater_approx.h b/src/tree/updater_approx.h
index 1e8f38da4269..2965663d8aa1 100644
--- a/src/tree/updater_approx.h
+++ b/src/tree/updater_approx.h
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2021 XGBoost contributors
+ * Copyright 2021-2022 XGBoost contributors
  *
  * \brief Implementation for the approx tree method.
  */
diff --git a/tests/cpp/gbm/test_gbtree.cc b/tests/cpp/gbm/test_gbtree.cc
index d703c5dbac4a..f9fe7d38660d 100644
--- a/tests/cpp/gbm/test_gbtree.cc
+++ b/tests/cpp/gbm/test_gbtree.cc
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2019-2021 XGBoost contributors
+ * Copyright 2019-2022 XGBoost contributors
  */
 #include <gtest/gtest.h>
 #include <dmlc/filesystem.h>
diff --git a/tests/cpp/predictor/test_cpu_predictor.cc b/tests/cpp/predictor/test_cpu_predictor.cc
index 279aacea54b7..f43747abdd9e 100644
--- a/tests/cpp/predictor/test_cpu_predictor.cc
+++ b/tests/cpp/predictor/test_cpu_predictor.cc
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2017-2020 XGBoost contributors
+ * Copyright 2017-2022 XGBoost contributors
  */
 #include <dmlc/filesystem.h>
 #include <gtest/gtest.h>
diff --git a/tests/cpp/tree/test_gpu_hist.cu b/tests/cpp/tree/test_gpu_hist.cu
index f1baa53eeead..2ae08a4e8bb3 100644
--- a/tests/cpp/tree/test_gpu_hist.cu
+++ b/tests/cpp/tree/test_gpu_hist.cu
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2017-2021 XGBoost contributors
+ * Copyright 2017-2022 XGBoost contributors
  */
 #include <gtest/gtest.h>
 #include <thrust/device_vector.h>

From 286f3a99f2ec6979327759235d04a8905db46774 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 19:41:06 +0800
Subject: [PATCH 097/124] Fix.

---
 src/common/partition_builder.h | 2 +-
 tests/cpp/tree/test_approx.cc  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index 0cf929a78c99..25debaf9db7d 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -294,7 +294,7 @@ class PartitionBuilder {
     size_t offset{0};
     h_nptr.push_back(offset);
     for (auto node : row_set) {
-      if (!tree[node.node_id].IsLeaf()) {
+      if (node.node_id < 0 || !tree[node.node_id].IsLeaf()) {
         continue;
       }
       if (node.begin) {  // guard for empty node.
diff --git a/tests/cpp/tree/test_approx.cc b/tests/cpp/tree/test_approx.cc
index 3aa29ddc623d..7d54e8ffc5a4 100644
--- a/tests/cpp/tree/test_approx.cc
+++ b/tests/cpp/tree/test_approx.cc
@@ -89,7 +89,7 @@ void TestLeafPartition(size_t n_samples) {
   }
 
   std::vector<size_t> h_nptr;
-  float split_value;
+  float split_value{0};
   for (auto const& page : Xy->GetBatches<GHistIndexMatrix>({Context::kCpuId, 64})) {
     bst_feature_t const split_ind = 0;
     auto ptr = page.cut.Ptrs()[split_ind + 1];

From c485d0a297b43bba122aecd504f57b2a34f64530 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Sat, 16 Apr 2022 22:42:09 +0800
Subject: [PATCH 098/124] Workaround.

---
 src/objective/regression_obj.cu        | 2 +-
 tests/python-gpu/test_gpu_with_dask.py | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 58cb60c235ea..5b24de17bd4f 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -696,7 +696,7 @@ void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_
   rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
   for (size_t i = 0; i < n_leaf; ++i) {
     if (n_valids[i] > 0) {
-      quantiles[i] /= static_cast<double>(n_valids[i]);
+      quantiles[i] /= static_cast<float>(n_valids[i]);
     } else {
       // Use original leaf value if no worker can provide the quantile.
       quantiles[i] = tree[h_node_idx[i]].LeafValue();
diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
index f4db033c81f0..2074ce073648 100644
--- a/tests/python-gpu/test_gpu_with_dask.py
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -201,9 +201,12 @@ def run_gpu_hist(
     )["history"]["train"][dataset.metric]
     note(history)
 
-    # See note on `ObjFunction::UpdateTreeLeaf.
+    # See note on `ObjFunction::UpdateTreeLeaf`.
     update_leaf = dataset.name.endswith("-l1")
-    if update_leaf and len(history) >= 2:
+    if update_leaf and len(history) == 2:
+        assert history[0] + 1e-2 >= history[-1]
+        return
+    if update_leaf and len(history) > 2:
         assert history[0] >= history[-1]
         return
     else:

From 56d8beb775f0e37867d8d548fa6573724ab87fef Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Sun, 17 Apr 2022 02:06:30 +0800
Subject: [PATCH 099/124] Fix set device.

---
 src/common/linalg_op.cuh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/common/linalg_op.cuh b/src/common/linalg_op.cuh
index 84caae2670cd..558a09ca6acb 100644
--- a/src/common/linalg_op.cuh
+++ b/src/common/linalg_op.cuh
@@ -13,6 +13,7 @@ namespace xgboost {
 namespace linalg {
 template <typename T, int32_t D, typename Fn>
 void ElementWiseKernelDevice(linalg::TensorView<T, D> t, Fn&& fn, cudaStream_t s = nullptr) {
+  dh::safe_cuda(cudaSetDevice(t.DeviceIdx()));
   static_assert(std::is_void<std::result_of_t<Fn(size_t, T&)>>::value,
                 "For function with return, use transform instead.");
   if (t.Contiguous()) {

From 02e04b6d150aaecb5c7b6c6371134cfce2b225d6 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sun, 17 Apr 2022 22:01:06 +0800
Subject: [PATCH 100/124] Small cleanup.

---
 src/tree/gpu_hist/row_partitioner.cuh | 4 ++--
 tests/python/test_with_dask.py        | 3 +--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 2ffb40322467..915b406bc4c7 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -13,12 +13,12 @@
 namespace xgboost {
 namespace tree {
 namespace detail {
-inline void FillMissingLeaf(std::vector<bst_node_t> const& missing, RowIndexCache* p_row_indices) {
+inline void FillMissingLeaf(std::vector<bst_node_t> const& maybe_missing, RowIndexCache* p_row_indices) {
   auto& row_indices = *p_row_indices;
   auto& h_node_idx = row_indices.node_idx.HostVector();
   auto& h_node_ptr = row_indices.node_ptr.HostVector();
 
-  for (auto leaf : missing) {
+  for (auto leaf : maybe_missing) {
     if (std::binary_search(h_node_idx.cbegin(), h_node_idx.cend(), leaf)) {
       continue;
     }
diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index 8566edb172b8..cbee3d72b254 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -1289,7 +1289,7 @@ def is_stump():
         def minimum_bin():
             return "max_bin" in params and params["max_bin"] == 2
 
-        # See note on `ObjFunction::UpdateTreeLeaf.
+        # See note on `ObjFunction::UpdateTreeLeaf`.
         update_leaf = dataset.name.endswith("-l1")
         if update_leaf and len(history) >= 2:
             assert history[0] >= history[-1]
@@ -1317,7 +1317,6 @@ def test_approx(
         self, client: "Client", params: Dict, dataset: tm.TestDataset
     ) -> None:
         num_rounds = 30
-        # params["eta"] = 0.1
         self.run_updater_test(client, params, num_rounds, dataset, 'approx')
 
     def run_quantile(self, name: str) -> None:

From 26dcc40debe2d21d1d4e5d48a8e9fc89f4d8b06d Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sun, 17 Apr 2022 22:06:23 +0800
Subject: [PATCH 101/124] Skip if updater doesn't support this feature.

---
 include/xgboost/tree_updater.h | 7 ++-----
 src/gbm/gbtree.cc              | 4 ++++
 src/gbm/gbtree.h               | 3 +++
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/include/xgboost/tree_updater.h b/include/xgboost/tree_updater.h
index 94264e3b6b8e..3268b84c8ff9 100644
--- a/include/xgboost/tree_updater.h
+++ b/include/xgboost/tree_updater.h
@@ -79,12 +79,9 @@ class TreeUpdater : public Configurable {
 
   virtual char const* Name() const = 0;
   /*!
-   * \brief Get the partition of rows based on the last iteration.
+   * \brief Get the partition of rows from the last iteration.
    */
-  virtual common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const {
-    LOG(FATAL) << "Objective is not supported by current tree method:" << this->Name();
-    return {};
-  }
+  virtual common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const { return {}; }
 
   /*!
    * \brief Create a tree updater given name
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index 2dade8720547..4f36357fa527 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -220,9 +220,13 @@ void CopyGradient(HostDeviceVector<GradientPair> const* in_gpair, int32_t n_thre
 void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const& predictions,
                             ObjFunction const* obj, size_t gidx,
                             std::vector<std::unique_ptr<RegTree>>* p_trees) {
+  CHECK(!updaters_.empty());
   if (!obj || !obj->Task().UpdateTreeLeaf()) {
     return;
   }
+  if (updaters_.back()->GetRowIndexCache(0).empty()) {
+    return;
+  }
   auto& trees = *p_trees;
   for (size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
     auto row_idx = updaters_.back()->GetRowIndexCache(tree_idx);
diff --git a/src/gbm/gbtree.h b/src/gbm/gbtree.h
index 4c4b7eb93756..0020c9a928dd 100644
--- a/src/gbm/gbtree.h
+++ b/src/gbm/gbtree.h
@@ -202,6 +202,9 @@ class GBTree : public GradientBooster {
   void ConfigureUpdaters();
   void ConfigureWithKnownData(Args const& cfg, DMatrix* fmat);
 
+  /**
+   * \brief Optionally update the leaf value.
+   */
   void UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const& predictions,
                       ObjFunction const* obj, size_t gidx,
                       std::vector<std::unique_ptr<RegTree>>* p_trees);

From 467348b48e8b3231ce7b08acebe714f0bc555d68 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 22 Apr 2022 18:32:23 +0800
Subject: [PATCH 102/124] Pass in position.

---
 include/xgboost/tree_updater.h          |  8 ++++----
 src/gbm/gbtree.cc                       | 18 +++++++++---------
 src/gbm/gbtree.h                        |  2 ++
 src/tree/updater_approx.cc              |  1 +
 src/tree/updater_colmaker.cc            |  6 +++---
 src/tree/updater_histmaker.cc           |  6 +++---
 src/tree/updater_prune.cc               |  8 ++++----
 src/tree/updater_quantile_hist.cc       |  1 +
 src/tree/updater_quantile_hist.h        |  4 ++--
 src/tree/updater_refresh.cc             |  6 +++---
 src/tree/updater_sync.cc                |  6 +++---
 tests/cpp/tree/test_histmaker.cc        |  6 ++++--
 tests/cpp/tree/test_prediction_cache.cc |  3 ++-
 tests/cpp/tree/test_prune.cc            | 11 ++++++-----
 tests/cpp/tree/test_refresh.cc          |  3 ++-
 tests/cpp/tree/test_tree_stat.cc        | 12 ++++++++----
 16 files changed, 57 insertions(+), 44 deletions(-)

diff --git a/include/xgboost/tree_updater.h b/include/xgboost/tree_updater.h
index 3268b84c8ff9..ccbbe9865674 100644
--- a/include/xgboost/tree_updater.h
+++ b/include/xgboost/tree_updater.h
@@ -53,14 +53,14 @@ class TreeUpdater : public Configurable {
    * \brief perform update to the tree models
    * \param gpair the gradient pair statistics of the data
    * \param data The data matrix passed to the updater.
-   * \param trees references the trees to be updated, updater will change the content of trees
+   * \param out_trees references the trees to be updated, updater will change the content of trees
    *   note: all the trees in the vector are updated, with the same statistics,
    *         but maybe different random seeds, usually one tree is passed in at a time,
    *         there can be multiple trees when we train random forest style model
    */
-  virtual void Update(HostDeviceVector<GradientPair>* gpair,
-                      DMatrix* data,
-                      const std::vector<RegTree*>& trees) = 0;
+  virtual void Update(HostDeviceVector<GradientPair>* gpair, DMatrix* data,
+                      common::Span<HostDeviceVector<bst_node_t>> out_position,
+                      const std::vector<RegTree*>& out_trees) = 0;
 
   /*!
    * \brief determines whether updater has enough knowledge about a given dataset
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index 4f36357fa527..35738100c9eb 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -236,7 +236,7 @@ void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const
 
 void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
                      PredictionCacheEntry* predt, ObjFunction const* obj) {
-  std::vector<std::vector<std::unique_ptr<RegTree> > > new_trees;
+  std::vector<std::vector<std::unique_ptr<RegTree>>> new_trees;
   const int ngroup = model_.learner_model_param->num_output_group;
   ConfigureWithKnownData(this->cfg_, p_fmat);
   monitor_.Start("BoostNewTrees");
@@ -328,10 +328,8 @@ void GBTree::InitUpdater(Args const& cfg) {
   }
 }
 
-void GBTree::BoostNewTrees(HostDeviceVector<GradientPair>* gpair,
-                           DMatrix *p_fmat,
-                           int bst_group,
-                           std::vector<std::unique_ptr<RegTree> >* ret) {
+void GBTree::BoostNewTrees(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat, int bst_group,
+                           std::vector<std::unique_ptr<RegTree>>* ret) {
   std::vector<RegTree*> new_trees;
   ret->clear();
   // create the trees
@@ -350,9 +348,9 @@ void GBTree::BoostNewTrees(HostDeviceVector<GradientPair>* gpair,
     } else if (tparam_.process_type == TreeProcessType::kUpdate) {
       for (auto const& up : updaters_) {
         CHECK(up->CanModifyTree())
-          << "Updater: `" << up->Name() << "` "
-          << "can not be used to modify existing trees. "
-          << "Set `process_type` to `default` if you want to build new trees.";
+            << "Updater: `" << up->Name() << "` "
+            << "can not be used to modify existing trees. "
+            << "Set `process_type` to `default` if you want to build new trees.";
       }
       CHECK_LT(model_.trees.size(), model_.trees_to_update.size())
           << "No more tree left for updating.  For updating existing trees, "
@@ -368,8 +366,10 @@ void GBTree::BoostNewTrees(HostDeviceVector<GradientPair>* gpair,
   CHECK_EQ(gpair->Size(), p_fmat->Info().num_row_)
       << "Mismatching size between number of rows from input data and size of "
          "gradient vector.";
+  node_position_.resize(new_trees.size());
   for (auto& up : updaters_) {
-    up->Update(gpair, p_fmat, new_trees);
+    up->Update(gpair, p_fmat, common::Span<HostDeviceVector<bst_node_t>>{node_position_},
+               new_trees);
   }
 }
 
diff --git a/src/gbm/gbtree.h b/src/gbm/gbtree.h
index 0020c9a928dd..885f9cd33db0 100644
--- a/src/gbm/gbtree.h
+++ b/src/gbm/gbtree.h
@@ -441,6 +441,8 @@ class GBTree : public GradientBooster {
   Args cfg_;
   // the updaters that can be applied to each of tree
   std::vector<std::unique_ptr<TreeUpdater>> updaters_;
+  // The node position for each row, 1 HDV for each tree in the forest.
+  std::vector<HostDeviceVector<bst_node_t>> node_position_;
   // Predictors
   std::unique_ptr<Predictor> cpu_predictor_;
 #if defined(XGBOOST_USE_CUDA)
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index a35518a07db7..c2819069d737 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -312,6 +312,7 @@ class GlobalApproxUpdater : public TreeUpdater {
   char const *Name() const override { return "grow_histmaker"; }
 
   void Update(HostDeviceVector<GradientPair> *gpair, DMatrix *m,
+              common::Span<HostDeviceVector<bst_node_t>> out_position,
               const std::vector<RegTree *> &trees) override {
     float lr = param_.learning_rate;
     param_.learning_rate = lr / trees.size();
diff --git a/src/tree/updater_colmaker.cc b/src/tree/updater_colmaker.cc
index e3d716f2cba8..6d63a00a139a 100644
--- a/src/tree/updater_colmaker.cc
+++ b/src/tree/updater_colmaker.cc
@@ -96,9 +96,9 @@ class ColMaker: public TreeUpdater {
     }
   }
 
-  void Update(HostDeviceVector<GradientPair> *gpair,
-              DMatrix* dmat,
-              const std::vector<RegTree*> &trees) override {
+  void Update(HostDeviceVector<GradientPair> *gpair, DMatrix *dmat,
+              common::Span<HostDeviceVector<bst_node_t>> out_position,
+              const std::vector<RegTree *> &trees) override {
     if (rabit::IsDistributed()) {
       LOG(FATAL) << "Updater `grow_colmaker` or `exact` tree method doesn't "
                     "support distributed training.";
diff --git a/src/tree/updater_histmaker.cc b/src/tree/updater_histmaker.cc
index 0a85d2d73832..27fc42455d2c 100644
--- a/src/tree/updater_histmaker.cc
+++ b/src/tree/updater_histmaker.cc
@@ -24,9 +24,9 @@ DMLC_REGISTRY_FILE_TAG(updater_histmaker);
 
 class HistMaker: public BaseMaker {
  public:
-  void Update(HostDeviceVector<GradientPair> *gpair,
-              DMatrix *p_fmat,
-              const std::vector<RegTree*> &trees) override {
+  void Update(HostDeviceVector<GradientPair> *gpair, DMatrix *p_fmat,
+              common::Span<HostDeviceVector<bst_node_t>> out_position,
+              const std::vector<RegTree *> &trees) override {
     interaction_constraints_.Configure(param_, p_fmat->Info().num_col_);
     // rescale learning rate according to size of trees
     float lr = param_.learning_rate;
diff --git a/src/tree/updater_prune.cc b/src/tree/updater_prune.cc
index f71f1c698cb9..dcda4a3b34a2 100644
--- a/src/tree/updater_prune.cc
+++ b/src/tree/updater_prune.cc
@@ -50,9 +50,9 @@ class TreePruner: public TreeUpdater {
   }
 
   // update the tree, do pruning
-  void Update(HostDeviceVector<GradientPair> *gpair,
-              DMatrix *p_fmat,
-              const std::vector<RegTree*> &trees) override {
+  void Update(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat,
+              common::Span<HostDeviceVector<bst_node_t>> out_position,
+              const std::vector<RegTree*>& trees) override {
     pruner_monitor_.Start("PrunerUpdate");
     // rescale learning rate according to size of trees
     float lr = param_.learning_rate;
@@ -61,7 +61,7 @@ class TreePruner: public TreeUpdater {
       this->DoPrune(tree);
     }
     param_.learning_rate = lr;
-    syncher_->Update(gpair, p_fmat, trees);
+    syncher_->Update(gpair, p_fmat, out_position, trees);
     pruner_monitor_.Stop("PrunerUpdate");
   }
 
diff --git a/src/tree/updater_quantile_hist.cc b/src/tree/updater_quantile_hist.cc
index b2f01e3b7917..de7247278cda 100644
--- a/src/tree/updater_quantile_hist.cc
+++ b/src/tree/updater_quantile_hist.cc
@@ -36,6 +36,7 @@ void QuantileHistMaker::Configure(const Args &args) {
 }
 
 void QuantileHistMaker::Update(HostDeviceVector<GradientPair> *gpair, DMatrix *dmat,
+                               common::Span<HostDeviceVector<bst_node_t>> out_position,
                                const std::vector<RegTree *> &trees) {
   // rescale learning rate according to size of trees
   float lr = param_.learning_rate;
diff --git a/src/tree/updater_quantile_hist.h b/src/tree/updater_quantile_hist.h
index cd7ff1dab626..dd6d41a02c9e 100644
--- a/src/tree/updater_quantile_hist.h
+++ b/src/tree/updater_quantile_hist.h
@@ -238,8 +238,8 @@ class QuantileHistMaker: public TreeUpdater {
   explicit QuantileHistMaker(ObjInfo task) : task_{task} {}
   void Configure(const Args& args) override;
 
-  void Update(HostDeviceVector<GradientPair>* gpair,
-              DMatrix* dmat,
+  void Update(HostDeviceVector<GradientPair>* gpair, DMatrix* dmat,
+              common::Span<HostDeviceVector<bst_node_t>> out_position,
               const std::vector<RegTree*>& trees) override;
 
   bool UpdatePredictionCache(const DMatrix *data,
diff --git a/src/tree/updater_refresh.cc b/src/tree/updater_refresh.cc
index d17c1e1444f7..8e82ae9f914c 100644
--- a/src/tree/updater_refresh.cc
+++ b/src/tree/updater_refresh.cc
@@ -42,9 +42,9 @@ class TreeRefresher: public TreeUpdater {
     return true;
   }
   // update the tree, do pruning
-  void Update(HostDeviceVector<GradientPair> *gpair,
-              DMatrix *p_fmat,
-              const std::vector<RegTree*> &trees) override {
+  void Update(HostDeviceVector<GradientPair> *gpair, DMatrix *p_fmat,
+              common::Span<HostDeviceVector<bst_node_t>> out_position,
+              const std::vector<RegTree *> &trees) override {
     if (trees.size() == 0) return;
     const std::vector<GradientPair> &gpair_h = gpair->ConstHostVector();
     // thread temporal space
diff --git a/src/tree/updater_sync.cc b/src/tree/updater_sync.cc
index 4f7c7a1a85a6..a4c1486fbf90 100644
--- a/src/tree/updater_sync.cc
+++ b/src/tree/updater_sync.cc
@@ -31,9 +31,9 @@ class TreeSyncher: public TreeUpdater {
     return "prune";
   }
 
-  void Update(HostDeviceVector<GradientPair>* ,
-              DMatrix*,
-              const std::vector<RegTree*> &trees) override {
+  void Update(HostDeviceVector<GradientPair>*, DMatrix*,
+              common::Span<HostDeviceVector<bst_node_t>> out_position,
+              const std::vector<RegTree*>& trees) override {
     if (rabit::GetWorldSize() == 1) return;
     std::string s_model;
     common::MemoryBufferStream fs(&s_model);
diff --git a/tests/cpp/tree/test_histmaker.cc b/tests/cpp/tree/test_histmaker.cc
index 56878b159d4b..90dc0a411294 100644
--- a/tests/cpp/tree/test_histmaker.cc
+++ b/tests/cpp/tree/test_histmaker.cc
@@ -39,7 +39,8 @@ TEST(GrowHistMaker, InteractionConstraint) {
     updater->Configure(Args{
         {"interaction_constraints", "[[0, 1]]"},
         {"num_feature", std::to_string(kCols)}});
-    updater->Update(&gradients, p_dmat.get(), {&tree});
+    std::vector<HostDeviceVector<bst_node_t>> position(1);
+    updater->Update(&gradients, p_dmat.get(), position, {&tree});
 
     ASSERT_EQ(tree.NumExtraNodes(), 4);
     ASSERT_EQ(tree[0].SplitIndex(), 1);
@@ -55,7 +56,8 @@ TEST(GrowHistMaker, InteractionConstraint) {
     std::unique_ptr<TreeUpdater> updater{
         TreeUpdater::Create("grow_histmaker", &param, ObjInfo{ObjInfo::kRegression})};
     updater->Configure(Args{{"num_feature", std::to_string(kCols)}});
-    updater->Update(&gradients, p_dmat.get(), {&tree});
+    std::vector<HostDeviceVector<bst_node_t>> position(1);
+    updater->Update(&gradients, p_dmat.get(), position, {&tree});
 
     ASSERT_EQ(tree.NumExtraNodes(), 10);
     ASSERT_EQ(tree[0].SplitIndex(), 1);
diff --git a/tests/cpp/tree/test_prediction_cache.cc b/tests/cpp/tree/test_prediction_cache.cc
index ebe66cf575b3..3e30e0699358 100644
--- a/tests/cpp/tree/test_prediction_cache.cc
+++ b/tests/cpp/tree/test_prediction_cache.cc
@@ -77,7 +77,8 @@ class TestPredictionCache : public ::testing::Test {
       std::vector<RegTree *> trees{&tree};
       auto gpair = GenerateRandomGradients(n_samples_);
       updater->Configure(Args{{"max_bin", "64"}});
-      updater->Update(&gpair, Xy_.get(), trees);
+      std::vector<HostDeviceVector<bst_node_t>> position(1);
+      updater->Update(&gpair, Xy_.get(), position, trees);
       HostDeviceVector<float> out_prediction_cached;
       out_prediction_cached.SetDevice(ctx.gpu_id);
       out_prediction_cached.Resize(n_samples_);
diff --git a/tests/cpp/tree/test_prune.cc b/tests/cpp/tree/test_prune.cc
index dc6a8da21d72..77f78b1399d9 100644
--- a/tests/cpp/tree/test_prune.cc
+++ b/tests/cpp/tree/test_prune.cc
@@ -43,22 +43,23 @@ TEST(Updater, Prune) {
   pruner->Configure(cfg);
 
   // loss_chg < min_split_loss;
+  std::vector<HostDeviceVector<bst_node_t>> position(trees.size());
   tree.ExpandNode(0, 0, 0, true, 0.0f, 0.3f, 0.4f, 0.0f, 0.0f,
                   /*left_sum=*/0.0f, /*right_sum=*/0.0f);
-  pruner->Update(&gpair, p_dmat.get(), trees);
+  pruner->Update(&gpair, p_dmat.get(), position, trees);
 
   ASSERT_EQ(tree.NumExtraNodes(), 0);
 
   // loss_chg > min_split_loss;
   tree.ExpandNode(0, 0, 0, true, 0.0f, 0.3f, 0.4f, 11.0f, 0.0f,
                   /*left_sum=*/0.0f, /*right_sum=*/0.0f);
-  pruner->Update(&gpair, p_dmat.get(), trees);
+  pruner->Update(&gpair, p_dmat.get(), position, trees);
 
   ASSERT_EQ(tree.NumExtraNodes(), 2);
 
   // loss_chg == min_split_loss;
   tree.Stat(0).loss_chg = 10;
-  pruner->Update(&gpair, p_dmat.get(), trees);
+  pruner->Update(&gpair, p_dmat.get(), position, trees);
 
   ASSERT_EQ(tree.NumExtraNodes(), 2);
 
@@ -74,7 +75,7 @@ TEST(Updater, Prune) {
                   /*left_sum=*/0.0f, /*right_sum=*/0.0f);
   cfg.emplace_back(std::make_pair("max_depth", "1"));
   pruner->Configure(cfg);
-  pruner->Update(&gpair, p_dmat.get(), trees);
+  pruner->Update(&gpair, p_dmat.get(), position, trees);
 
   ASSERT_EQ(tree.NumExtraNodes(), 2);
 
@@ -84,7 +85,7 @@ TEST(Updater, Prune) {
                   /*left_sum=*/0.0f, /*right_sum=*/0.0f);
   cfg.emplace_back(std::make_pair("min_split_loss", "0"));
   pruner->Configure(cfg);
-  pruner->Update(&gpair, p_dmat.get(), trees);
+  pruner->Update(&gpair, p_dmat.get(), position, trees);
   ASSERT_EQ(tree.NumExtraNodes(), 2);
 }
 }  // namespace tree
diff --git a/tests/cpp/tree/test_refresh.cc b/tests/cpp/tree/test_refresh.cc
index 5b71f0841e19..f0abd0a871aa 100644
--- a/tests/cpp/tree/test_refresh.cc
+++ b/tests/cpp/tree/test_refresh.cc
@@ -44,7 +44,8 @@ TEST(Updater, Refresh) {
   tree.Stat(cright).base_weight = 1.3;
 
   refresher->Configure(cfg);
-  refresher->Update(&gpair, p_dmat.get(), trees);
+  std::vector<HostDeviceVector<bst_node_t>> position;
+  refresher->Update(&gpair, p_dmat.get(), position, trees);
 
   bst_float constexpr kEps = 1e-6;
   ASSERT_NEAR(-0.183392, tree[cright].LeafValue(), kEps);
diff --git a/tests/cpp/tree/test_tree_stat.cc b/tests/cpp/tree/test_tree_stat.cc
index 772420ce0f23..723ca34ebc93 100644
--- a/tests/cpp/tree/test_tree_stat.cc
+++ b/tests/cpp/tree/test_tree_stat.cc
@@ -27,7 +27,8 @@ class UpdaterTreeStatTest : public ::testing::Test {
     up->Configure(Args{});
     RegTree tree;
     tree.param.num_feature = kCols;
-    up->Update(&gpairs_, p_dmat_.get(), {&tree});
+    std::vector<HostDeviceVector<bst_node_t>> position(1);
+    up->Update(&gpairs_, p_dmat_.get(), position, {&tree});
 
     tree.WalkTree([&tree](bst_node_t nidx) {
       if (tree[nidx].IsLeaf()) {
@@ -87,13 +88,15 @@ class UpdaterEtaTest : public ::testing::Test {
       RegTree tree_0;
       {
         tree_0.param.num_feature = kCols;
-        up_0->Update(&gpairs_, p_dmat_.get(), {&tree_0});
+        std::vector<HostDeviceVector<bst_node_t>> position(1);
+        up_0->Update(&gpairs_, p_dmat_.get(), position, {&tree_0});
       }
 
       RegTree tree_1;
       {
         tree_1.param.num_feature = kCols;
-        up_1->Update(&gpairs_, p_dmat_.get(), {&tree_1});
+        std::vector<HostDeviceVector<bst_node_t>> position(1);
+        up_1->Update(&gpairs_, p_dmat_.get(), position, {&tree_1});
       }
       tree_0.WalkTree([&](bst_node_t nidx) {
         if (tree_0[nidx].IsLeaf()) {
@@ -149,7 +152,8 @@ class TestMinSplitLoss : public ::testing::Test {
     up->Configure(args);
 
     RegTree tree;
-    up->Update(&gpair_, dmat_.get(), {&tree});
+    std::vector<HostDeviceVector<bst_node_t>> position(1);
+    up->Update(&gpair_, dmat_.get(), position, {&tree});
 
     auto n_nodes = tree.NumExtraNodes();
     return n_nodes;

From 423520c3daa6253c7abce2338c406a7044da3403 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 22 Apr 2022 19:01:00 +0800
Subject: [PATCH 103/124] Copy position on host.

---
 include/xgboost/objective.h                |  2 +-
 include/xgboost/tree_updater.h             |  4 ----
 src/common/partition_builder.h             | 19 ++++---------------
 src/gbm/gbtree.cc                          |  7 ++-----
 src/objective/regression_obj.cu            |  4 ++--
 src/tree/updater_approx.cc                 | 22 ++++++++--------------
 src/tree/updater_approx.h                  |  4 ++--
 src/tree/updater_quantile_hist.cc          | 15 ++++++++-------
 src/tree/updater_quantile_hist.h           | 16 +++++-----------
 tests/cpp/objective/test_regression_obj.cc |  4 ++--
 tests/cpp/tree/test_approx.cc              | 20 ++++++++++----------
 11 files changed, 44 insertions(+), 73 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 1fd3515c1cca..84da73e9d71d 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -103,7 +103,7 @@ class ObjFunction : public Configurable {
    * \param info MetaInfo providing labels and weights.
    * \param p_tree Tree that needs to be updated.
    */
-  virtual void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
+  virtual void UpdateTreeLeaf(HostDeviceVector<bst_node_t> const& row_index, MetaInfo const& info,
                               HostDeviceVector<float> const& prediction, RegTree* p_tree) const {}
 
   /*!
diff --git a/include/xgboost/tree_updater.h b/include/xgboost/tree_updater.h
index ccbbe9865674..a5352aa0a4cf 100644
--- a/include/xgboost/tree_updater.h
+++ b/include/xgboost/tree_updater.h
@@ -78,10 +78,6 @@ class TreeUpdater : public Configurable {
   }
 
   virtual char const* Name() const = 0;
-  /*!
-   * \brief Get the partition of rows from the last iteration.
-   */
-  virtual common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const { return {}; }
 
   /*!
    * \brief Create a tree updater given name
diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index 25debaf9db7d..06973979c723 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -283,16 +283,11 @@ class PartitionBuilder {
   // Copy row partitions into global cache for reuse in objective
   template <typename Sampledp>
   void LeafPartition(Context const* ctx, RegTree const& tree, RowSetCollection const& row_set,
-                     std::vector<RowIndexCache>* p_out_row_indices, Sampledp sampledp) const {
-    p_out_row_indices->emplace_back(ctx, row_set.Data()->size());
-    auto& h_ridx = p_out_row_indices->back().row_index.HostVector();
-    auto& h_nptr = p_out_row_indices->back().node_ptr.HostVector();
-    auto& h_nidx = p_out_row_indices->back().node_idx.HostVector();
-    CHECK(h_nptr.empty());
+                     std::vector<bst_node_t>* p_position, Sampledp sampledp) const {
+    auto& h_pos = *p_position;
+    h_pos.resize(row_set.Data()->size(), std::numeric_limits<bst_node_t>::max());
 
     auto p_begin = row_set.Data()->data();
-    size_t offset{0};
-    h_nptr.push_back(offset);
     for (auto node : row_set) {
       if (node.node_id < 0 || !tree[node.node_id].IsLeaf()) {
         continue;
@@ -301,16 +296,10 @@ class PartitionBuilder {
         size_t ptr_offset = node.end - p_begin;
         CHECK_LE(ptr_offset, row_set.Data()->size()) << node.node_id;
         for (auto idx = node.begin; idx != node.end; ++idx) {
-          if (!sampledp(*idx)) {
-            h_ridx[offset++] = *idx;
-          }
+          h_pos[*idx] = sampledp(*idx) ? ~node.node_id : node.node_id;
         }
       }
-      h_nptr.push_back(offset);
-      h_nidx.push_back(node.node_id);
     }
-    CHECK_LE(offset, row_set.Data()->size());
-    h_ridx.resize(offset);
   }
 
  protected:
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index 35738100c9eb..a96dccfc405b 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -224,13 +224,10 @@ void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const
   if (!obj || !obj->Task().UpdateTreeLeaf()) {
     return;
   }
-  if (updaters_.back()->GetRowIndexCache(0).empty()) {
-    return;
-  }
   auto& trees = *p_trees;
   for (size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
-    auto row_idx = updaters_.back()->GetRowIndexCache(tree_idx);
-    obj->UpdateTreeLeaf(row_idx, p_fmat->Info(), predictions, trees[tree_idx].get());
+    auto const& position = this->node_position_.at(tree_idx);
+    obj->UpdateTreeLeaf(position, p_fmat->Info(), predictions, trees[tree_idx].get());
   }
 }
 
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 5b24de17bd4f..6d1d365206b1 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -852,10 +852,10 @@ class MeanAbsoluteError : public ObjFunction {
     });
   }
 
-  void UpdateTreeLeaf(common::Span<RowIndexCache const> row_index, MetaInfo const& info,
+  void UpdateTreeLeaf(HostDeviceVector<bst_node_t> const& position, MetaInfo const& info,
                       HostDeviceVector<float> const& prediction, RegTree* p_tree) const override {
     if (ctx_->IsCPU()) {
-      detail::UpdateTreeLeafHost(ctx_, row_index, info, prediction, 0.5, p_tree);
+      // detail::UpdateTreeLeafHost(ctx_, row_index, info, prediction, 0.5, p_tree);
     } else {
 #if defined(XGBOOST_USE_CUDA)
       detail::UpdateTreeLeafDevice(ctx_, row_index, info, prediction, 0.5, p_tree);
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index c2819069d737..823603df4860 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -156,7 +156,7 @@ class GloablApproxBuilder {
   }
 
   void LeafPartition(RegTree const &tree, common::Span<float> hess,
-                     std::vector<RowIndexCache> *p_out_row_indices) {
+                     std::vector<bst_node_t> *p_out_row_indices) {
     monitor_->Start(__func__);
     if (!evaluator_.Task().UpdateTreeLeaf()) {
       return;
@@ -179,7 +179,7 @@ class GloablApproxBuilder {
         monitor_{monitor} {}
 
   void UpdateTree(DMatrix *p_fmat, std::vector<GradientPair> const &gpair, common::Span<float> hess,
-                  RegTree *p_tree, std::vector<RowIndexCache> *p_out_row_indices) {
+                  RegTree *p_tree, HostDeviceVector<bst_node_t> *p_out_row_indices) {
     p_last_tree_ = p_tree;
     this->InitData(p_fmat, hess);
 
@@ -246,7 +246,8 @@ class GloablApproxBuilder {
       expand_set = driver.Pop();
     }
 
-    this->LeafPartition(tree, hess, p_out_row_indices);
+    auto &h_row_indices = p_out_row_indices->HostVector();
+    this->LeafPartition(tree, hess, &h_row_indices);
   }
 };
 
@@ -265,8 +266,6 @@ class GlobalApproxUpdater : public TreeUpdater {
   DMatrix *cached_{nullptr};
   std::shared_ptr<common::ColumnSampler> column_sampler_ =
       std::make_shared<common::ColumnSampler>();
-  // cache for row partitions
-  std::vector<std::vector<RowIndexCache>> row_set_collection_;
   ObjInfo task_;
 
  public:
@@ -293,7 +292,6 @@ class GlobalApproxUpdater : public TreeUpdater {
     sampled->resize(h_gpair.size());
     std::copy(h_gpair.cbegin(), h_gpair.cend(), sampled->begin());
     auto &rnd = common::GlobalRandom();
-    row_set_collection_.clear();
 
     if (param.subsample != 1.0) {
       CHECK(param.sampling_method != TrainParam::kGradientBased)
@@ -334,14 +332,14 @@ class GlobalApproxUpdater : public TreeUpdater {
 
     cached_ = m;
 
+    size_t t_idx = 0;
     for (auto p_tree : trees) {
-      row_set_collection_.emplace_back();
-      auto &row_indices = row_set_collection_.back();
       if (hist_param_.single_precision_histogram) {
-        this->f32_impl_->UpdateTree(m, h_gpair, hess, p_tree, &row_indices);
+        this->f32_impl_->UpdateTree(m, h_gpair, hess, p_tree, &out_position[t_idx]);
       } else {
-        this->f64_impl_->UpdateTree(m, h_gpair, hess, p_tree, &row_indices);
+        this->f64_impl_->UpdateTree(m, h_gpair, hess, p_tree, &out_position[t_idx]);
       }
+      ++t_idx;
     }
     param_.learning_rate = lr;
   }
@@ -358,10 +356,6 @@ class GlobalApproxUpdater : public TreeUpdater {
     }
     return true;
   }
-
-  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
-    return row_set_collection_.at(tree_idx);
-  }
 };
 
 DMLC_REGISTRY_FILE_TAG(grow_histmaker);
diff --git a/src/tree/updater_approx.h b/src/tree/updater_approx.h
index 2965663d8aa1..bb37f99ec61d 100644
--- a/src/tree/updater_approx.h
+++ b/src/tree/updater_approx.h
@@ -124,8 +124,8 @@ class ApproxRowPartitioner {
   auto const &Partitions() const { return row_set_collection_; }
 
   void LeafPartition(Context const *ctx, RegTree const &tree, common::Span<float const> hess,
-                     std::vector<RowIndexCache> *p_out_row_indices) const {
-    partition_builder_.LeafPartition(ctx, tree, this->Partitions(), p_out_row_indices,
+                     std::vector<bst_node_t> *p_out_position) const {
+    partition_builder_.LeafPartition(ctx, tree, this->Partitions(), p_out_position,
                                      [&](size_t idx) -> bool { return hess[idx] - .0f == .0f; });
   }
 
diff --git a/src/tree/updater_quantile_hist.cc b/src/tree/updater_quantile_hist.cc
index de7247278cda..9fbb0b37ba5e 100644
--- a/src/tree/updater_quantile_hist.cc
+++ b/src/tree/updater_quantile_hist.cc
@@ -54,15 +54,15 @@ void QuantileHistMaker::Update(HostDeviceVector<GradientPair> *gpair, DMatrix *d
     }
   }
 
-  row_set_collection_.clear();
+  size_t t_idx{0};
   for (auto p_tree : trees) {
-    row_set_collection_.emplace_back();
-    auto &row_indices = row_set_collection_.back();
+    auto &row_indices = out_position[t_idx];
     if (hist_maker_param_.single_precision_histogram) {
       this->float_builder_->UpdateTree(gpair, dmat, p_tree, &row_indices);
     } else {
       this->double_builder_->UpdateTree(gpair, dmat, p_tree, &row_indices);
     }
+    ++t_idx;
   }
 
   param_.learning_rate = lr;
@@ -176,7 +176,7 @@ void QuantileHistMaker::Builder<GradientSumT>::BuildHistogram(
 template <typename GradientSumT>
 void QuantileHistMaker::Builder<GradientSumT>::LeafPartition(
     RegTree const &tree, common::Span<GradientPair const> gpair,
-    std::vector<RowIndexCache> *p_out_row_indices) {
+    std::vector<bst_node_t> *p_out_row_indices) {
   monitor_->Start(__func__);
   if (!evaluator_->Task().UpdateTreeLeaf()) {
     return;
@@ -191,7 +191,7 @@ void QuantileHistMaker::Builder<GradientSumT>::LeafPartition(
 template <typename GradientSumT>
 void QuantileHistMaker::Builder<GradientSumT>::ExpandTree(
     DMatrix *p_fmat, RegTree *p_tree, const std::vector<GradientPair> &gpair_h,
-    std::vector<RowIndexCache> *p_out_row_indices) {
+    HostDeviceVector<bst_node_t> *p_out_row_indices) {
   monitor_->Start(__func__);
 
   Driver<CPUExpandEntry> driver(static_cast<TrainParam::TreeGrowPolicy>(param_.grow_policy));
@@ -248,14 +248,15 @@ void QuantileHistMaker::Builder<GradientSumT>::ExpandTree(
     expand_set = driver.Pop();
   }
 
-  this->LeafPartition(tree, gpair_h, p_out_row_indices);
+  auto &h_row_indices = p_out_row_indices->HostVector();
+  this->LeafPartition(tree, gpair_h, &h_row_indices);
   monitor_->Stop(__func__);
 }
 
 template <typename GradientSumT>
 void QuantileHistMaker::Builder<GradientSumT>::UpdateTree(
     HostDeviceVector<GradientPair> *gpair, DMatrix *p_fmat, RegTree *p_tree,
-    std::vector<RowIndexCache> *p_out_row_indices) {
+    HostDeviceVector<bst_node_t> *p_out_row_indices) {
   monitor_->Start(__func__);
 
   std::vector<GradientPair> *gpair_ptr = &(gpair->HostVector());
diff --git a/src/tree/updater_quantile_hist.h b/src/tree/updater_quantile_hist.h
index dd6d41a02c9e..ec81f23ac5e1 100644
--- a/src/tree/updater_quantile_hist.h
+++ b/src/tree/updater_quantile_hist.h
@@ -218,9 +218,9 @@ class HistRowPartitioner {
 
   void LeafPartition(Context const* ctx, RegTree const& tree,
                      common::Span<GradientPair const> gpair,
-                     std::vector<RowIndexCache>* p_out_row_indices) const {
+                     std::vector<bst_node_t>* p_out_position) const {
     partition_builder_.LeafPartition(
-        ctx, tree, this->Partitions(), p_out_row_indices,
+        ctx, tree, this->Partitions(), p_out_position,
         [&](size_t idx) -> bool { return gpair[idx].GetHess() - .0f == .0f; });
   }
 
@@ -276,10 +276,6 @@ class QuantileHistMaker: public TreeUpdater {
     return "grow_quantile_histmaker";
   }
 
-  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
-    return row_set_collection_.at(tree_idx);
-  }
-
  protected:
   CPUHistMakerTrainParam hist_maker_param_;
   // training parameter
@@ -304,7 +300,7 @@ class QuantileHistMaker: public TreeUpdater {
     }
     // update one tree, growing
     void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat, RegTree* p_tree,
-                    std::vector<RowIndexCache>* p_out_row_indices);
+                    HostDeviceVector<bst_node_t>* p_out_row_indices);
 
     bool UpdatePredictionCache(DMatrix const* data, linalg::VectorView<float> out_preds) const;
 
@@ -324,10 +320,10 @@ class QuantileHistMaker: public TreeUpdater {
                         std::vector<GradientPair> const& gpair);
 
     void LeafPartition(RegTree const& tree, common::Span<GradientPair const> gpair,
-                       std::vector<RowIndexCache>* p_out_row_indices);
+                       std::vector<bst_node_t>* p_out_row_indices);
 
     void ExpandTree(DMatrix* p_fmat, RegTree* p_tree, const std::vector<GradientPair>& gpair_h,
-                    std::vector<RowIndexCache>* p_out_row_indices);
+                    HostDeviceVector<bst_node_t>* p_out_row_indices);
 
    private:
     const size_t n_trees_;
@@ -353,8 +349,6 @@ class QuantileHistMaker: public TreeUpdater {
   };
 
  protected:
-  // cache for row partitions
-  std::vector<std::vector<RowIndexCache>> row_set_collection_;
   std::unique_ptr<Builder<float>> float_builder_;
   std::unique_ptr<Builder<double>> double_builder_;
   ObjInfo task_;
diff --git a/tests/cpp/objective/test_regression_obj.cc b/tests/cpp/objective/test_regression_obj.cc
index 96242b8d798e..9e260c54a508 100644
--- a/tests/cpp/objective/test_regression_obj.cc
+++ b/tests/cpp/objective/test_regression_obj.cc
@@ -419,7 +419,7 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteError)) {
   for (size_t i = 0; i < h_predt.size(); ++i) {
     h_predt[i] = labels[i] + i;
   }
-  obj->UpdateTreeLeaf(common::Span<RowIndexCache const>{row_idx}, info, predt, &tree);
+  // obj->UpdateTreeLeaf(common::Span<RowIndexCache const>{row_idx}, info, predt, &tree);
   ASSERT_EQ(tree[1].LeafValue(), -1);
   ASSERT_EQ(tree[2].LeafValue(), -4);
 }
@@ -458,7 +458,7 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteErrorLeaf)) {
   ASSERT_EQ(tree.GetNumLeaves(), 4);
 
   auto empty_leaf = tree[4].LeafValue();
-  obj->UpdateTreeLeaf(row_idx_v, info, predt, &tree);
+  // obj->UpdateTreeLeaf(row_idx_v, info, predt, &tree);
   ASSERT_EQ(tree[3].LeafValue(), -5);
   ASSERT_EQ(tree[4].LeafValue(), empty_leaf);
   ASSERT_EQ(tree[5].LeafValue(), -10);
diff --git a/tests/cpp/tree/test_approx.cc b/tests/cpp/tree/test_approx.cc
index 7d54e8ffc5a4..966eebb2ba98 100644
--- a/tests/cpp/tree/test_approx.cc
+++ b/tests/cpp/tree/test_approx.cc
@@ -96,18 +96,18 @@ void TestLeafPartition(size_t n_samples) {
     split_value = page.cut.Values().at(ptr / 2);
     GetSplit(&tree, split_value, &candidates);
     partitioner.UpdatePosition(&ctx, page, candidates, &tree);
-    std::vector<RowIndexCache> cache;
+    std::vector<bst_node_t> cache;
     partitioner.LeafPartition(&ctx, tree, hess, &cache);
-    auto const& row_idx = cache.front();
-    ASSERT_EQ(n, row_idx.row_index.Size());
-    h_nptr = row_idx.node_ptr.ConstHostVector();
-    ASSERT_EQ(h_nptr.size(), 3);
-    ASSERT_EQ(h_nptr[0], 0);
-    ASSERT_EQ(h_nptr[2], n);  // equal to sampled rows
+    // auto const& row_idx = cache.front();
+    // ASSERT_EQ(n, row_idx.row_index.Size());
+    // h_nptr = row_idx.node_ptr.ConstHostVector();
+    // ASSERT_EQ(h_nptr.size(), 3);
+    // ASSERT_EQ(h_nptr[0], 0);
+    // ASSERT_EQ(h_nptr[2], n);  // equal to sampled rows
 
-    ASSERT_EQ(row_idx.node_idx.Size(), 2);
-    ASSERT_EQ(row_idx.node_idx.HostVector()[0], 1);
-    ASSERT_EQ(row_idx.node_idx.HostVector()[1], 2);
+    // ASSERT_EQ(row_idx.node_idx.Size(), 2);
+    // ASSERT_EQ(row_idx.node_idx.HostVector()[0], 1);
+    // ASSERT_EQ(row_idx.node_idx.HostVector()[1], 2);
   }
 
   for (auto const& page : Xy->GetBatches<SparsePage>()) {

From 4ac121165bd1575fc2ff4e863a14a62ebfdb2d41 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 22 Apr 2022 19:59:59 +0800
Subject: [PATCH 104/124] Compile.

---
 src/common/common.h             | 14 ++++++++++
 src/data/data.cc                | 11 +-------
 src/objective/regression_obj.cu | 46 +++++++++++++++++++++------------
 3 files changed, 45 insertions(+), 26 deletions(-)

diff --git a/src/common/common.h b/src/common/common.h
index d909a37691ae..d281ed414ef3 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -268,6 +268,20 @@ template <typename Indexable>
 XGBOOST_DEVICE size_t LastOf(size_t group, Indexable const &indptr) {
   return indptr[group + 1] - 1;
 }
+
+template <typename T, typename Idx>
+void RunLengthEncode(std::vector<T> const &sorted_values, std::vector<Idx> *p_out) {
+  auto &out = *p_out;
+  out = std::vector<Idx>{0};
+  for (size_t i = 1; i < sorted_values.size(); ++i) {
+    if (sorted_values[i] != sorted_values[i - 1]) {
+      out.push_back(i);
+    }
+  }
+  if (out.back() != sorted_values.size()) {
+    out.push_back(sorted_values.size());
+  }
+}
 }  // namespace common
 }  // namespace xgboost
 #endif  // XGBOOST_COMMON_COMMON_H_
diff --git a/src/data/data.cc b/src/data/data.cc
index 86f73523a39d..0aec340e7582 100644
--- a/src/data/data.cc
+++ b/src/data/data.cc
@@ -512,16 +512,7 @@ void MetaInfo::SetInfoFromHost(Context const& ctx, StringView key, Json arr) {
       }
     }
     CHECK(non_dec) << "`qid` must be sorted in non-decreasing order along with data.";
-    group_ptr_.clear();
-    group_ptr_.push_back(0);
-    for (size_t i = 1; i < query_ids.size(); ++i) {
-      if (query_ids[i] != query_ids[i - 1]) {
-        group_ptr_.push_back(i);
-      }
-    }
-    if (group_ptr_.back() != query_ids.size()) {
-      group_ptr_.push_back(query_ids.size());
-    }
+    common::RunLengthEncode(query_ids, &group_ptr_);
     data::ValidateQueryGroup(group_ptr_);
     return;
   }
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 6d1d365206b1..f9621a8cd8c7 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -671,11 +671,11 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .set_body([]() { return new TweedieRegression(); });
 
 namespace detail {
-void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_index,
-                      RegTree* p_tree) {
+void UpdateLeafValues(std::vector<float>* p_quantiles, std::vector<size_t> const& row_index,
+                      std::vector<bst_node_t> const nidx, RegTree* p_tree) {
   auto& tree = *p_tree;
   auto& quantiles = *p_quantiles;
-  auto const& h_node_idx = row_index.node_idx.HostVector();
+  auto const& h_node_idx = nidx;
 
   size_t n_leaf{h_node_idx.size()};
   rabit::Allreduce<rabit::op::Max>(&n_leaf, 1);
@@ -703,7 +703,7 @@ void UpdateLeafValues(std::vector<float>* p_quantiles, RowIndexCache const& row_
     }
   }
 
-  for (size_t i = 0; i < row_index.node_idx.Size(); ++i) {
+  for (size_t i = 0; i < nidx.size(); ++i) {
     auto nidx = h_node_idx[i];
     auto q = quantiles[i];
     CHECK(tree[nidx].IsLeaf());
@@ -753,18 +753,31 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
 }
 #endif  // defined(XGBOOST_USE_CUDA)
 
-void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> row_index,
+void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& position,
                         MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
                         RegTree* p_tree) {
   auto& tree = *p_tree;
-  CHECK(!row_index.empty());
-  std::vector<float> quantiles(row_index.front().node_idx.Size(), 0);
-  std::vector<int32_t> n_valids(quantiles.size(), 0);
-  // loop over external memory partitions
-  for (auto const& part : row_index) {
-    std::vector<float> results(part.node_idx.Size());
-    auto const& h_node_idx = part.node_idx.ConstHostVector();
-    auto const& h_node_ptr = part.node_ptr.ConstHostVector();
+  CHECK(!position.empty());
+
+  auto ridx = common::ArgSort<size_t>(position);
+  std::vector<bst_node_t> nidx(position);
+  // permute
+  for (size_t i = 0; i < position.size(); ++i) {
+    nidx[i] = position[ridx[i]];
+  }
+  std::vector<size_t> segments{0};
+  common::RunLengthEncode(nidx, &segments);
+  CHECK_GT(segments.size(), 0);
+  size_t n_leaf = segments.size() - 1;
+  nidx.resize(n_leaf);
+
+  std::vector<float> quantiles(n_leaf, 0);
+  std::vector<int32_t> n_valids(n_leaf, 0);
+
+  {
+    std::vector<float> results(nidx.size());
+    auto const& h_node_idx = nidx;
+    auto const& h_node_ptr = segments;
     CHECK_LE(h_node_ptr.back(), info.num_row_);
     // loop over each leaf
     common::ParallelFor(results.size(), ctx->Threads(), [&](size_t k) {
@@ -772,7 +785,7 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
       CHECK(tree[nidx].IsLeaf());
       CHECK_LT(k + 1, h_node_ptr.size());
       size_t n = h_node_ptr[k + 1] - h_node_ptr[k];
-      auto h_row_set = part.row_index.HostSpan().subspan(h_node_ptr[k], n);
+      auto h_row_set = common::Span<size_t const>{ridx}.subspan(h_node_ptr[k], n);
       // multi-target not yet supported.
       auto h_labels = info.labels.HostView().Slice(linalg::All(), 0);
       auto const& h_predt = predt.ConstHostVector();
@@ -817,7 +830,7 @@ void UpdateTreeLeafHost(Context const* ctx, common::Span<RowIndexCache const> ro
     }
   }
 
-  UpdateLeafValues(&quantiles, row_index.front(), p_tree);
+  UpdateLeafValues(&quantiles, ridx, nidx, p_tree);
 }
 }  // namespace detail
 
@@ -855,7 +868,8 @@ class MeanAbsoluteError : public ObjFunction {
   void UpdateTreeLeaf(HostDeviceVector<bst_node_t> const& position, MetaInfo const& info,
                       HostDeviceVector<float> const& prediction, RegTree* p_tree) const override {
     if (ctx_->IsCPU()) {
-      // detail::UpdateTreeLeafHost(ctx_, row_index, info, prediction, 0.5, p_tree);
+      auto const& h_position = position.ConstHostVector();
+      detail::UpdateTreeLeafHost(ctx_, h_position, info, prediction, 0.5, p_tree);
     } else {
 #if defined(XGBOOST_USE_CUDA)
       detail::UpdateTreeLeafDevice(ctx_, row_index, info, prediction, 0.5, p_tree);

From b66c2627ba543a926548bdcfd39dc3a204f83a7e Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 22 Apr 2022 20:05:59 +0800
Subject: [PATCH 105/124] Fix n leaf.

---
 src/objective/regression_obj.cu   | 2 ++
 src/tree/updater_approx.cc        | 1 -
 src/tree/updater_quantile_hist.cc | 1 -
 3 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index f9621a8cd8c7..8644f4feaa87 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -769,6 +769,8 @@ void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& posit
   common::RunLengthEncode(nidx, &segments);
   CHECK_GT(segments.size(), 0);
   size_t n_leaf = segments.size() - 1;
+  auto n_unique = std::unique(nidx.begin(), nidx.end()) - nidx.begin();
+  CHECK_EQ(n_unique, n_leaf);
   nidx.resize(n_leaf);
 
   std::vector<float> quantiles(n_leaf, 0);
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index 823603df4860..05b0c2abaf45 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -161,7 +161,6 @@ class GloablApproxBuilder {
     if (!evaluator_.Task().UpdateTreeLeaf()) {
       return;
     }
-    CHECK(p_out_row_indices->empty());
     for (auto const &part : partitioner_) {
       part.LeafPartition(ctx_, tree, hess, p_out_row_indices);
     }
diff --git a/src/tree/updater_quantile_hist.cc b/src/tree/updater_quantile_hist.cc
index 9fbb0b37ba5e..e870166217f0 100644
--- a/src/tree/updater_quantile_hist.cc
+++ b/src/tree/updater_quantile_hist.cc
@@ -181,7 +181,6 @@ void QuantileHistMaker::Builder<GradientSumT>::LeafPartition(
   if (!evaluator_->Task().UpdateTreeLeaf()) {
     return;
   }
-  CHECK(p_out_row_indices->empty());
   for (auto const &part : partitioner_) {
     part.LeafPartition(ctx_, tree, gpair, p_out_row_indices);
   }

From 6b4f6f463012fbb6798b4118977cbf88159fbed6 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 22 Apr 2022 21:51:49 +0800
Subject: [PATCH 106/124] Sampling.

---
 include/xgboost/objective.h                |  2 +-
 src/common/common.h                        | 16 ++++---
 src/data/data.cc                           |  2 +-
 src/objective/regression_obj.cu            | 29 +++++++++---
 src/tree/updater_approx.cc                 |  6 +--
 tests/cpp/objective/test_regression_obj.cc | 51 +++++++++++-----------
 6 files changed, 63 insertions(+), 43 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 84da73e9d71d..52b87c5f2944 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -103,7 +103,7 @@ class ObjFunction : public Configurable {
    * \param info MetaInfo providing labels and weights.
    * \param p_tree Tree that needs to be updated.
    */
-  virtual void UpdateTreeLeaf(HostDeviceVector<bst_node_t> const& row_index, MetaInfo const& info,
+  virtual void UpdateTreeLeaf(HostDeviceVector<bst_node_t> const& position, MetaInfo const& info,
                               HostDeviceVector<float> const& prediction, RegTree* p_tree) const {}
 
   /*!
diff --git a/src/common/common.h b/src/common/common.h
index d281ed414ef3..d36d4a80ac0e 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -269,17 +269,21 @@ XGBOOST_DEVICE size_t LastOf(size_t group, Indexable const &indptr) {
   return indptr[group + 1] - 1;
 }
 
-template <typename T, typename Idx>
-void RunLengthEncode(std::vector<T> const &sorted_values, std::vector<Idx> *p_out) {
+/**
+ * \brief Run length encode on CPU, input must be sorted.
+ */
+template <typename Iter, typename Idx>
+void RunLengthEncode(Iter begin, Iter end, std::vector<Idx> *p_out) {
   auto &out = *p_out;
   out = std::vector<Idx>{0};
-  for (size_t i = 1; i < sorted_values.size(); ++i) {
-    if (sorted_values[i] != sorted_values[i - 1]) {
+  size_t n = std::distance(begin, end);
+  for (size_t i = 1; i < n; ++i) {
+    if (begin[i] != begin[i - 1]) {
       out.push_back(i);
     }
   }
-  if (out.back() != sorted_values.size()) {
-    out.push_back(sorted_values.size());
+  if (out.back() != n) {
+    out.push_back(n);
   }
 }
 }  // namespace common
diff --git a/src/data/data.cc b/src/data/data.cc
index 0aec340e7582..c297527c6bae 100644
--- a/src/data/data.cc
+++ b/src/data/data.cc
@@ -512,7 +512,7 @@ void MetaInfo::SetInfoFromHost(Context const& ctx, StringView key, Json arr) {
       }
     }
     CHECK(non_dec) << "`qid` must be sorted in non-decreasing order along with data.";
-    common::RunLengthEncode(query_ids, &group_ptr_);
+    common::RunLengthEncode(query_ids.cbegin(), query_ids.cend(), &group_ptr_);
     data::ValidateQueryGroup(group_ptr_);
     return;
   }
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 8644f4feaa87..1c4ee2c6b01e 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -760,18 +760,33 @@ void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& posit
   CHECK(!position.empty());
 
   auto ridx = common::ArgSort<size_t>(position);
-  std::vector<bst_node_t> nidx(position);
-  // permute
+  std::vector<bst_node_t> sorted_pos(position);
+  // permutation
   for (size_t i = 0; i < position.size(); ++i) {
-    nidx[i] = position[ridx[i]];
+    sorted_pos[i] = position[ridx[i]];
   }
-  std::vector<size_t> segments{0};
-  common::RunLengthEncode(nidx, &segments);
+  // find the first non-sampled row
+  auto begin_pos =
+      std::distance(sorted_pos.cbegin(), std::find_if(sorted_pos.cbegin(), sorted_pos.cend(),
+                                                      [](bst_node_t nidx) { return nidx >= 0; }));
+  CHECK_LE(begin_pos, sorted_pos.size());
+  if (begin_pos == sorted_pos.size()) {
+    return;
+  }
+
+  std::vector<size_t> segments;
+  auto beg_it = sorted_pos.begin() + begin_pos;
+  common::RunLengthEncode(beg_it, sorted_pos.end(), &segments);
   CHECK_GT(segments.size(), 0);
+  // skip the sampled rows in indptr
+  std::transform(segments.begin(), segments.end(), segments.begin(),
+                 [begin_pos](size_t ptr) { return ptr + begin_pos; });
+
   size_t n_leaf = segments.size() - 1;
-  auto n_unique = std::unique(nidx.begin(), nidx.end()) - nidx.begin();
+  auto n_unique = std::unique(beg_it, sorted_pos.end()) - beg_it;
   CHECK_EQ(n_unique, n_leaf);
-  nidx.resize(n_leaf);
+  std::vector<bst_node_t> nidx(n_leaf);
+  std::copy(beg_it, beg_it + n_unique, nidx.begin());
 
   std::vector<float> quantiles(n_leaf, 0);
   std::vector<int32_t> n_valids(n_leaf, 0);
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index 05b0c2abaf45..bf668b2a5e46 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -178,7 +178,7 @@ class GloablApproxBuilder {
         monitor_{monitor} {}
 
   void UpdateTree(DMatrix *p_fmat, std::vector<GradientPair> const &gpair, common::Span<float> hess,
-                  RegTree *p_tree, HostDeviceVector<bst_node_t> *p_out_row_indices) {
+                  RegTree *p_tree, HostDeviceVector<bst_node_t> *p_out_position) {
     p_last_tree_ = p_tree;
     this->InitData(p_fmat, hess);
 
@@ -245,8 +245,8 @@ class GloablApproxBuilder {
       expand_set = driver.Pop();
     }
 
-    auto &h_row_indices = p_out_row_indices->HostVector();
-    this->LeafPartition(tree, hess, &h_row_indices);
+    auto &h_position = p_out_position->HostVector();
+    this->LeafPartition(tree, hess, &h_position);
   }
 };
 
diff --git a/tests/cpp/objective/test_regression_obj.cc b/tests/cpp/objective/test_regression_obj.cc
index 9e260c54a508..964bb92e41c9 100644
--- a/tests/cpp/objective/test_regression_obj.cc
+++ b/tests/cpp/objective/test_regression_obj.cc
@@ -398,20 +398,14 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteError)) {
 
   RegTree tree;
   tree.ExpandNode(0, /*split_index=*/1, 2, true, 0.0f, 2.f, 3.f, 4.f, 2.f, 1.f, 1.f);
-  std::vector<RowIndexCache> row_idx;
-  row_idx.emplace_back(&ctx, info.labels.Shape(0));
-
-  row_idx.back().node_idx.HostVector().push_back(1);  // left
-  row_idx.back().node_idx.HostVector().push_back(2);  // right
-  auto& ptr = row_idx.back().node_ptr.HostVector();
-  ptr.push_back(0);
-  ptr.push_back(3);
-  ptr.push_back(info.labels.Size());
-  auto& h_row_idx = row_idx.back().row_index.HostVector();
-  for (size_t i = info.labels.Size() - 1;; --i) {
-    h_row_idx[i] = i;
-    if (i == 0) {
-      break;
+
+  HostDeviceVector<bst_node_t> position(labels.size(), 0);
+  auto& h_position = position.HostVector();
+  for (size_t i = 0; i < labels.size(); ++i) {
+    if (i < labels.size() / 2) {
+      h_position[i] = 1;  // left
+    } else {
+      h_position[i] = 2;  // right
     }
   }
 
@@ -419,7 +413,8 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteError)) {
   for (size_t i = 0; i < h_predt.size(); ++i) {
     h_predt[i] = labels[i] + i;
   }
-  // obj->UpdateTreeLeaf(common::Span<RowIndexCache const>{row_idx}, info, predt, &tree);
+
+  obj->UpdateTreeLeaf(position, info, predt, &tree);
   ASSERT_EQ(tree[1].LeafValue(), -1);
   ASSERT_EQ(tree[2].LeafValue(), -4);
 }
@@ -441,15 +436,21 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteErrorLeaf)) {
     h_predt[i] = h_labels[i] + i;
   }
 
-  std::vector<RowIndexCache> row_idx_v;
-  row_idx_v.emplace_back(&ctx, info.labels.Shape(0));
-
-  auto& part = row_idx_v.back();
-  part.node_idx = {3, 4, 5, 6};
-  // starting from 3 to emulate subsampling, empty leaaft for node 4.
-  part.node_ptr = {3, 8, 8, 13, 16};
-  auto& h_row_idx = part.row_index.HostVector();
-  std::iota(h_row_idx.begin(), h_row_idx.end(), 0);
+  HostDeviceVector<bst_node_t> position(info.labels.Size(), 0);
+  auto& h_position = position.HostVector();
+  for (int32_t i = 0; i < 3; ++i) {
+    h_position[i] = ~i;  // negation for sampled nodes.
+  }
+  for (size_t i = 3; i < 8; ++i) {
+    h_position[i] = 3;
+  }
+  // empty leaf for node 4
+  for (size_t i = 8; i < 13; ++i) {
+    h_position[i] = 5;
+  }
+  for (size_t i = 13; i < h_labels.size(); ++i) {
+    h_position[i] = 6;
+  }
 
   RegTree tree;
   tree.ExpandNode(0, /*split_index=*/1, 2, true, 0.0f, 2.f, 3.f, 4.f, 2.f, 1.f, 1.f);
@@ -458,7 +459,7 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteErrorLeaf)) {
   ASSERT_EQ(tree.GetNumLeaves(), 4);
 
   auto empty_leaf = tree[4].LeafValue();
-  // obj->UpdateTreeLeaf(row_idx_v, info, predt, &tree);
+  obj->UpdateTreeLeaf(position, info, predt, &tree);
   ASSERT_EQ(tree[3].LeafValue(), -5);
   ASSERT_EQ(tree[4].LeafValue(), empty_leaf);
   ASSERT_EQ(tree[5].LeafValue(), -10);

From 2cebd450bd052af72e607e6fa1efab0dcc69f0ee Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Fri, 22 Apr 2022 22:23:58 +0800
Subject: [PATCH 107/124] approx test.

---
 tests/cpp/tree/test_approx.cc | 35 ++++++++++++++++++++++-------------
 1 file changed, 22 insertions(+), 13 deletions(-)

diff --git a/tests/cpp/tree/test_approx.cc b/tests/cpp/tree/test_approx.cc
index 966eebb2ba98..2e2fd4a0b3d7 100644
--- a/tests/cpp/tree/test_approx.cc
+++ b/tests/cpp/tree/test_approx.cc
@@ -26,7 +26,7 @@ TEST(Approx, Partitioner) {
   std::transform(grad.HostVector().cbegin(), grad.HostVector().cend(), hess.begin(),
                  [](auto gpair) { return gpair.GetHess(); });
 
-  for (auto const &page : Xy->GetBatches<GHistIndexMatrix>({64, hess, true})) {
+  for (auto const& page : Xy->GetBatches<GHistIndexMatrix>({64, hess, true})) {
     bst_feature_t const split_ind = 0;
     {
       auto min_value = page.cut.MinValues()[split_ind];
@@ -96,20 +96,29 @@ void TestLeafPartition(size_t n_samples) {
     split_value = page.cut.Values().at(ptr / 2);
     GetSplit(&tree, split_value, &candidates);
     partitioner.UpdatePosition(&ctx, page, candidates, &tree);
-    std::vector<bst_node_t> cache;
-    partitioner.LeafPartition(&ctx, tree, hess, &cache);
-    // auto const& row_idx = cache.front();
-    // ASSERT_EQ(n, row_idx.row_index.Size());
-    // h_nptr = row_idx.node_ptr.ConstHostVector();
-    // ASSERT_EQ(h_nptr.size(), 3);
-    // ASSERT_EQ(h_nptr[0], 0);
-    // ASSERT_EQ(h_nptr[2], n);  // equal to sampled rows
+    std::vector<bst_node_t> position;
+    partitioner.LeafPartition(&ctx, tree, hess, &position);
+    std::sort(position.begin(), position.end());
+    size_t beg = std::distance(
+        position.begin(),
+        std::find_if(position.begin(), position.end(), [&](bst_node_t nidx) { return nidx >= 0; }));
+    std::vector<size_t> nptr;
+    common::RunLengthEncode(position.cbegin() + beg, position.cend(), &nptr);
+    std::transform(nptr.begin(), nptr.end(), nptr.begin(), [&](size_t x) { return x + beg; });
+    auto n_uniques = std::unique(position.begin() + beg, position.end()) - (position.begin() + beg);
+    ASSERT_EQ(nptr.size(), n_uniques + 1);
+    ASSERT_EQ(nptr[0], beg);
+    ASSERT_EQ(nptr.back(), n_samples);
 
-    // ASSERT_EQ(row_idx.node_idx.Size(), 2);
-    // ASSERT_EQ(row_idx.node_idx.HostVector()[0], 1);
-    // ASSERT_EQ(row_idx.node_idx.HostVector()[1], 2);
+    h_nptr = nptr;
   }
 
+  if (h_nptr.front() == n_samples) {
+    return;
+  }
+
+  ASSERT_GE(h_nptr.size(), 2);
+
   for (auto const& page : Xy->GetBatches<SparsePage>()) {
     auto batch = page.GetView();
     size_t left{0};
@@ -118,7 +127,7 @@ void TestLeafPartition(size_t n_samples) {
         left++;
       }
     }
-    ASSERT_EQ(left, h_nptr[1]);  // equal to number of sampled assigned to left
+    ASSERT_EQ(left, h_nptr[1] - h_nptr[0]);  // equal to number of sampled assigned to left
   }
 }
 }  // anonymous namespace

From 13f772bac60747dc4c527db8d421e4a6e4e72679 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 22 Apr 2022 23:53:57 +0800
Subject: [PATCH 108/124] Compile on GPU.

---
 src/common/device_helpers.cuh                 |   1 +
 src/objective/adaptive.cuh                    | 156 ++++++++++++++++++
 src/objective/regression_obj.cu               |  46 ++++--
 src/tree/gpu_hist/row_partitioner.cuh         | 139 +---------------
 src/tree/updater_gpu_hist.cu                  |  46 ++----
 .../cpp/objective/test_regression_obj_gpu.cu  |  38 ++++-
 .../cpp/tree/gpu_hist/test_row_partitioner.cu |  59 +------
 tests/cpp/tree/test_gpu_hist.cu               |   3 +-
 8 files changed, 255 insertions(+), 233 deletions(-)
 create mode 100644 src/objective/adaptive.cuh

diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
index 2286db60d8de..334e3b4f89bf 100644
--- a/src/common/device_helpers.cuh
+++ b/src/common/device_helpers.cuh
@@ -1637,5 +1637,6 @@ class CUDAStream {
   }
 
   CUDAStreamView View() const { return CUDAStreamView{stream_}; }
+  void Sync() { this->View().Sync(); }
 };
 }  // namespace dh
diff --git a/src/objective/adaptive.cuh b/src/objective/adaptive.cuh
new file mode 100644
index 000000000000..e5a32fc246f6
--- /dev/null
+++ b/src/objective/adaptive.cuh
@@ -0,0 +1,156 @@
+/*!
+ * Copyright 2022 by XGBoost Contributors
+ */
+#pragma once
+#include <thrust/sort.h>
+
+#include <cub/cub.cuh>
+
+#include "../common/device_helpers.cuh"
+#include "xgboost/generic_parameters.h"
+#include "xgboost/host_device_vector.h"
+#include "xgboost/tree_model.h"
+
+namespace xgboost {
+namespace obj {
+namespace detail {
+inline void FillMissingLeaf(std::vector<bst_node_t> const& maybe_missing,
+                            HostDeviceVector<bst_node_t>* p_nidx,
+                            HostDeviceVector<size_t>* p_nptr) {
+  auto& h_node_idx = p_nidx->HostVector();
+  auto& h_node_ptr = p_nptr->HostVector();
+
+  for (auto leaf : maybe_missing) {
+    if (std::binary_search(h_node_idx.cbegin(), h_node_idx.cend(), leaf)) {
+      continue;
+    }
+    auto it = std::upper_bound(h_node_idx.cbegin(), h_node_idx.cend(), leaf);
+    auto pos = it - h_node_idx.cbegin();
+    h_node_idx.insert(h_node_idx.cbegin() + pos, leaf);
+    h_node_ptr.insert(h_node_ptr.cbegin() + pos, h_node_ptr[pos]);
+  }
+
+  // push to device.
+  p_nidx->ConstDevicePointer();
+  p_nptr->ConstDevicePointer();
+}
+
+inline void EncodeTreeLeaf(Context const* ctx, common::Span<bst_node_t const> position,
+                           HostDeviceVector<size_t>* p_nptr, HostDeviceVector<bst_node_t>* p_nidx,
+                           RegTree const& tree) {
+  // copy position to buffer
+  dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
+  size_t n_samples = position.size();
+  dh::XGBDeviceAllocator<char> alloc;
+  dh::device_vector<bst_node_t> sorted_position(position.size());
+  dh::safe_cuda(cudaMemcpyAsync(sorted_position.data().get(), position.data(),
+                                position.size_bytes(), cudaMemcpyDeviceToDevice));
+  dh::device_vector<size_t> ridx(position.size());
+  dh::Iota(dh::ToSpan(ridx));
+  // sort row index according to node index
+  thrust::stable_sort_by_key(thrust::cuda::par(alloc), sorted_position.begin(),
+                             sorted_position.begin() + n_samples, ridx.begin());
+
+  size_t n_leaf = tree.GetNumLeaves();
+  // +1 for subsample, which is set to an unique value in above kernel.
+  size_t max_n_unique = n_leaf + 1;
+
+  dh::caching_device_vector<size_t> counts_out(max_n_unique + 1, 0);
+  auto d_counts_out = dh::ToSpan(counts_out).subspan(0, max_n_unique);
+  auto d_num_runs_out = dh::ToSpan(counts_out).subspan(max_n_unique, 1);
+  dh::caching_device_vector<bst_node_t> unique_out(max_n_unique, 0);
+  auto d_unique_out = dh::ToSpan(unique_out);
+
+  size_t nbytes;
+  cub::DeviceRunLengthEncode::Encode(nullptr, nbytes, sorted_position.begin(),
+                                     unique_out.data().get(), counts_out.data().get(),
+                                     d_num_runs_out.data(), n_samples);
+  dh::TemporaryArray<char> temp(nbytes);
+  cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, sorted_position.begin(),
+                                     unique_out.data().get(), counts_out.data().get(),
+                                     d_num_runs_out.data(), n_samples);
+
+  dh::XGBCachingDeviceAllocator<char> caching;
+  dh::PinnedMemory pinned_pool;
+  auto pinned = pinned_pool.GetSpan<char>(sizeof(size_t) + sizeof(bst_node_t));
+  dh::CUDAStream copy_stream;
+  size_t* h_num_runs = reinterpret_cast<size_t*>(pinned.subspan(0, sizeof(size_t)).data());
+  // flag for whether there's ignored position
+  bst_node_t* h_first_unique =
+      reinterpret_cast<bst_node_t*>(pinned.subspan(sizeof(size_t), sizeof(bst_node_t)).data());
+  dh::safe_cuda(cudaMemcpyAsync(h_num_runs, d_num_runs_out.data(), sizeof(size_t),
+                                cudaMemcpyDeviceToHost, copy_stream.View()));
+  dh::safe_cuda(cudaMemcpyAsync(h_first_unique, d_unique_out.data(), sizeof(bst_node_t),
+                                cudaMemcpyDeviceToHost, copy_stream.View()));
+
+  /**
+   * copy node index (leaf index)
+   */
+  auto& nidx = *p_nidx;
+  auto& nptr = *p_nptr;
+  nidx.SetDevice(ctx->gpu_id);
+  nidx.Resize(n_leaf);
+  auto d_node_idx = nidx.DeviceSpan();
+
+  nptr.SetDevice(ctx->gpu_id);
+  nptr.Resize(n_leaf + 1, 0);
+  auto d_node_ptr = nptr.DeviceSpan();
+
+  dh::LaunchN(n_leaf, [=] XGBOOST_DEVICE(size_t i) {
+    if (i >= d_num_runs_out[0]) {
+      // d_num_runs_out <= max_n_unique
+      // this omits all the leaf that are empty. A leaf can be empty when there's
+      // missing data, which can be caused by sparse input and distributed training.
+      return;
+    }
+    if (d_unique_out[0] < 0) {
+      // shift 1 to the left
+      // some rows are ignored due to sampling, `kIgnoredTreePosition` is -1 so it's the
+      // smallest value and is sorted to the left.
+      // d_unique_out.size() == n_leaf + 1.
+      d_node_idx[i] = d_unique_out[i + 1];
+      d_node_ptr[i + 1] = d_counts_out[i + 1];
+      if (i == 0) {
+        d_node_ptr[0] = d_counts_out[0];
+      }
+    } else {
+      d_node_idx[i] = d_unique_out[i];
+      d_node_ptr[i + 1] = d_counts_out[i];
+      if (i == 0) {
+        d_node_ptr[0] = 0;
+      }
+    }
+  });
+  thrust::inclusive_scan(thrust::cuda::par(caching), dh::tbegin(d_node_ptr), dh::tend(d_node_ptr),
+                         dh::tbegin(d_node_ptr));
+  copy_stream.View().Sync();
+  if (*h_first_unique < 0) {
+    *h_num_runs -= 1;  // sampled.
+  }
+  CHECK_GT(*h_num_runs, 0);
+  CHECK_LE(*h_num_runs, n_leaf);
+
+  if (*h_num_runs < n_leaf) {
+    // shrink to omit the `kIgnoredTreePosition`.
+    nptr.Resize(*h_num_runs + 1);
+    nidx.Resize(*h_num_runs);
+
+    std::vector<bst_node_t> leaves;
+    tree.WalkTree([&](bst_node_t nidx) {
+      if (tree[nidx].IsLeaf()) {
+        leaves.push_back(nidx);
+      }
+      return true;
+    });
+    CHECK_EQ(leaves.size(), n_leaf);
+    // Fill all the leaves that don't have any sample. This is hacky and inefficient. An
+    // alternative is to leave the objective to handle missing leaf, which is more messy
+    // as we need to take other distributed workers into account.
+    FillMissingLeaf(leaves, &nidx, &nptr);
+  }
+  CHECK_EQ(nidx.Size(), n_leaf);
+  CHECK_EQ(nptr.Size(), n_leaf + 1);
+}
+}  // namespace detail
+}  // namespace obj
+}  // namespace xgboost
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 1c4ee2c6b01e..a7260965cdec 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -32,9 +32,10 @@
 #include "xgboost/span.h"
 
 #if defined(XGBOOST_USE_CUDA)
-#include "../common/linalg_op.cuh"
 #include "../common/device_helpers.cuh"
+#include "../common/linalg_op.cuh"
 #include "../common/stats.cuh"
+#include "adaptive.cuh"
 #endif  // defined(XGBOOST_USE_CUDA)
 
 namespace xgboost {
@@ -671,8 +672,8 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .set_body([]() { return new TweedieRegression(); });
 
 namespace detail {
-void UpdateLeafValues(std::vector<float>* p_quantiles, std::vector<size_t> const& row_index,
-                      std::vector<bst_node_t> const nidx, RegTree* p_tree) {
+void UpdateLeafValues(std::vector<float>* p_quantiles, std::vector<bst_node_t> const nidx,
+                      RegTree* p_tree) {
   auto& tree = *p_tree;
   auto& quantiles = *p_quantiles;
   auto const& h_node_idx = nidx;
@@ -712,24 +713,27 @@ void UpdateLeafValues(std::vector<float>* p_quantiles, std::vector<size_t> const
 }
 
 #if defined(XGBOOST_USE_CUDA)
-void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const> row_index,
+
+
+void UpdateTreeLeafDevice(Context const* ctx, common::Span<bst_node_t const> position,
                           MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
                           RegTree* p_tree) {
   dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
-  CHECK_EQ(row_index.size(), 1)
-      << "External memory with GPU hist should have only 1 row partition.";
-  auto const& part = row_index.front();
+  dh::device_vector<size_t> ridx;
+  HostDeviceVector<size_t> nptr;
+  HostDeviceVector<bst_node_t> nidx;
+
+  EncodeTreeLeaf(ctx, position, &nptr, &nidx, *p_tree);
 
   HostDeviceVector<float> quantiles;
   predt.SetDevice(ctx->gpu_id);
   auto d_predt = predt.ConstDeviceSpan();
   auto d_labels = info.labels.View(ctx->gpu_id);
 
-  part.row_index.SetDevice(ctx->gpu_id);
-  auto d_row_index = part.row_index.ConstDeviceSpan();
-  part.node_ptr.SetDevice(ctx->gpu_id);
-  auto seg_beg = part.node_ptr.ConstDeviceSpan().data();
-  auto seg_end = seg_beg + part.node_ptr.Size();
+
+  auto d_row_index = dh::ToSpan(ridx);
+  auto seg_beg = nptr.DevicePointer();
+  auto seg_end = seg_beg + nptr.Size();
   auto val_beg = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
                                                   [=] XGBOOST_DEVICE(size_t i) {
                                                     auto predt = d_predt[d_row_index[i]];
@@ -737,7 +741,7 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
                                                     return y - predt;
                                                   });
   auto val_end = val_beg + d_labels.Size();
-  CHECK_EQ(part.node_idx.Size() + 1, part.node_ptr.Size());
+  CHECK_EQ(nidx.Size() + 1, nptr.Size());
   if (info.weights_.Empty()) {
     common::SegmentedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, &quantiles);
   } else {
@@ -749,7 +753,7 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<RowIndexCache const>
                                       w_it + d_weights.size(), &quantiles);
   }
 
-  UpdateLeafValues(&quantiles.HostVector(), row_index.front(), p_tree);
+  UpdateLeafValues(&quantiles.HostVector(), nidx.ConstHostVector(), p_tree);
 }
 #endif  // defined(XGBOOST_USE_CUDA)
 
@@ -771,6 +775,15 @@ void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& posit
                                                       [](bst_node_t nidx) { return nidx >= 0; }));
   CHECK_LE(begin_pos, sorted_pos.size());
   if (begin_pos == sorted_pos.size()) {
+    std::vector<bst_node_t> leaf;
+    tree.WalkTree([&](bst_node_t nidx) {
+      if (tree[nidx].IsLeaf()) {
+        leaf.push_back(nidx);
+      }
+      return true;
+    });
+    std::vector<float> quantiles;
+    UpdateLeafValues(&quantiles, leaf, p_tree);
     return;
   }
 
@@ -847,7 +860,7 @@ void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& posit
     }
   }
 
-  UpdateLeafValues(&quantiles, ridx, nidx, p_tree);
+  UpdateLeafValues(&quantiles, nidx, p_tree);
 }
 }  // namespace detail
 
@@ -889,7 +902,8 @@ class MeanAbsoluteError : public ObjFunction {
       detail::UpdateTreeLeafHost(ctx_, h_position, info, prediction, 0.5, p_tree);
     } else {
 #if defined(XGBOOST_USE_CUDA)
-      detail::UpdateTreeLeafDevice(ctx_, row_index, info, prediction, 0.5, p_tree);
+      auto d_position = position.ConstDeviceSpan();
+      detail::UpdateTreeLeafDevice(ctx_, d_position, info, prediction, 0.5, p_tree);
 #else
       common::AssertGPUSupport();
 #endif  //  defined(XGBOOST_USE_CUDA)
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 915b406bc4c7..52d3ca578e45 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -12,28 +12,6 @@
 
 namespace xgboost {
 namespace tree {
-namespace detail {
-inline void FillMissingLeaf(std::vector<bst_node_t> const& maybe_missing, RowIndexCache* p_row_indices) {
-  auto& row_indices = *p_row_indices;
-  auto& h_node_idx = row_indices.node_idx.HostVector();
-  auto& h_node_ptr = row_indices.node_ptr.HostVector();
-
-  for (auto leaf : maybe_missing) {
-    if (std::binary_search(h_node_idx.cbegin(), h_node_idx.cend(), leaf)) {
-      continue;
-    }
-    auto it = std::upper_bound(h_node_idx.cbegin(), h_node_idx.cend(), leaf);
-    auto pos = it - h_node_idx.cbegin();
-    h_node_idx.insert(h_node_idx.cbegin() + pos, leaf);
-    h_node_ptr.insert(h_node_ptr.cbegin() + pos, h_node_ptr[pos]);
-  }
-
-  // push to device.
-  row_indices.node_idx.ConstDevicePointer();
-  row_indices.node_ptr.ConstDevicePointer();
-}
-}  // namespace detail
-
 /*! \brief Count how many rows are assigned to left node. */
 __forceinline__ __device__ void AtomicIncrement(int64_t* d_count, bool increment) {
 #if __CUDACC_VER_MAJOR__ > 8
@@ -188,7 +166,7 @@ class RowPartitioner {
    */
   template <typename FinalisePositionOpT, typename Sampledp>
   void FinalisePosition(Context const* ctx, RegTree const* p_tree, size_t n_leaf, ObjInfo task,
-                        std::vector<RowIndexCache>* p_out_row_indices, FinalisePositionOpT op,
+                        HostDeviceVector<bst_node_t>* p_out_row_indices, FinalisePositionOpT op,
                         Sampledp sampledp) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
@@ -205,124 +183,19 @@ class RowPartitioner {
       return;
     }
 
-    auto sorted_position = position_.Other();
+    p_out_row_indices->SetDevice(ctx->gpu_id);
+    p_out_row_indices->Resize(position_.Size());
+    auto sorted_position = p_out_row_indices->DevicePointer();
     dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
       auto position = d_position[idx];
       RowIndexT ridx = d_ridx[idx];
       bst_node_t new_position = op(ridx, position);
-      sorted_position[ridx] = sampledp(ridx) ? kIgnoredTreePosition : new_position;
-      if (new_position == kIgnoredTreePosition) {
+      sorted_position[ridx] = sampledp(ridx) ? ~new_position : new_position;
+      if (new_position < 0) {
         return;
       }
       d_position[idx] = new_position;
     });
-
-    // copy position to buffer
-    size_t n_samples = position_.Size();
-    dh::XGBDeviceAllocator<char> alloc;
-    auto& row_indices = p_out_row_indices->back();
-    // sort row index according to node index
-    row_indices.row_index.SetDevice(ctx->gpu_id);
-    row_indices.row_index.Resize(ridx_.Size());
-    dh::Iota(row_indices.row_index.DeviceSpan());
-    thrust::stable_sort_by_key(thrust::cuda::par(alloc), sorted_position,
-                               sorted_position + n_samples, row_indices.row_index.DevicePointer());
-
-    // +1 for subsample, which is set to an unique value in above kernel.
-    size_t max_n_unique = n_leaf + 1;
-
-    dh::caching_device_vector<size_t> counts_out(max_n_unique + 1, 0);
-    auto d_counts_out = dh::ToSpan(counts_out).subspan(0, max_n_unique);
-    auto d_num_runs_out = dh::ToSpan(counts_out).subspan(max_n_unique, 1);
-    dh::caching_device_vector<bst_node_t> unique_out(max_n_unique, 0);
-    auto d_unique_out = dh::ToSpan(unique_out);
-
-    size_t nbytes;
-    cub::DeviceRunLengthEncode::Encode(nullptr, nbytes, sorted_position, unique_out.data().get(),
-                                       counts_out.data().get(), d_num_runs_out.data(), n_samples);
-    dh::TemporaryArray<char> temp(nbytes);
-    cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, sorted_position,
-                                       unique_out.data().get(), counts_out.data().get(),
-                                       d_num_runs_out.data(), n_samples);
-
-    dh::XGBCachingDeviceAllocator<char> caching;
-    auto pinned = pinned_.GetSpan<char>(sizeof(size_t) + sizeof(bst_node_t));
-    size_t* h_num_runs = reinterpret_cast<size_t*>(pinned.subspan(0, sizeof(size_t)).data());
-    // flag for whether there's ignored position
-    bst_node_t* h_first_unique =
-        reinterpret_cast<bst_node_t*>(pinned.subspan(sizeof(size_t), sizeof(bst_node_t)).data());
-    dh::safe_cuda(cudaMemcpyAsync(h_num_runs, d_num_runs_out.data(), sizeof(size_t),
-                                  cudaMemcpyDeviceToHost, streams_[0]));
-    dh::safe_cuda(cudaMemcpyAsync(h_first_unique, d_unique_out.data(), sizeof(bst_node_t),
-                                  cudaMemcpyDeviceToHost, streams_[0]));
-
-    /**
-     * copy node index (leaf index)
-     */
-    row_indices.node_idx.SetDevice(ctx->gpu_id);
-    row_indices.node_idx.Resize(n_leaf);
-    auto d_node_idx = row_indices.node_idx.DeviceSpan();
-
-    row_indices.node_ptr.SetDevice(ctx->gpu_id);
-    row_indices.node_ptr.Resize(n_leaf + 1, 0);
-    auto d_node_ptr = row_indices.node_ptr.DeviceSpan();
-
-    dh::LaunchN(n_leaf, [=] XGBOOST_DEVICE(size_t i) {
-      if (i >= d_num_runs_out[0]) {
-        // d_num_runs_out <= max_n_unique
-        // this omits all the leaf that are empty. A leaf can be empty when there's
-        // missing data, which can be caused by sparse input and distributed training.
-        return;
-      }
-      if (d_unique_out[0] == kIgnoredTreePosition) {
-        // shift 1 to the left
-        // some rows are ignored due to sampling, `kIgnoredTreePosition` is -1 so it's the
-        // smallest value and is sorted to the left.
-        // d_unique_out.size() == n_leaf + 1.
-        d_node_idx[i] = d_unique_out[i + 1];
-        d_node_ptr[i + 1] = d_counts_out[i + 1];
-        if (i == 0) {
-          d_node_ptr[0] = d_counts_out[0];
-        }
-      } else {
-        d_node_idx[i] = d_unique_out[i];
-        d_node_ptr[i + 1] = d_counts_out[i];
-        if (i == 0) {
-          d_node_ptr[0] = 0;
-        }
-      }
-    });
-    thrust::inclusive_scan(thrust::cuda::par(caching), dh::tbegin(d_node_ptr), dh::tend(d_node_ptr),
-                           dh::tbegin(d_node_ptr));
-    dh::CUDAStreamView{streams_[0]}.Sync();
-    if (*h_first_unique == kIgnoredTreePosition) {
-      *h_num_runs -= 1;  // sampled.
-    }
-    CHECK_GT(*h_num_runs, 0);
-    CHECK_LE(*h_num_runs, n_leaf);
-
-    if (*h_num_runs < n_leaf) {
-      // empty leaf, have to fill in all missing leaves
-      auto const& tree = *p_tree;
-      // shrink to omit the `kIgnoredTreePosition`.
-      row_indices.node_ptr.Resize(*h_num_runs + 1);
-      row_indices.node_idx.Resize(*h_num_runs);
-
-      std::vector<bst_node_t> leaves;
-      tree.WalkTree([&](bst_node_t nidx) {
-        if (tree[nidx].IsLeaf()) {
-          leaves.push_back(nidx);
-        }
-        return true;
-      });
-      CHECK_EQ(leaves.size(), n_leaf);
-      // Fill all the leaves that don't have any sample. This is hacky and inefficient. An
-      // alternative is to leave the objective to handle missing leaf, which is more messy
-      // as we need to take other distributed workers into account.
-      detail::FillMissingLeaf(leaves, &row_indices);
-    }
-    CHECK_EQ(row_indices.node_idx.Size(), n_leaf);
-    CHECK_EQ(row_indices.node_ptr.Size(), n_leaf + 1);
   }
 
   /**
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 6c287fb85eab..0bb6c601db80 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -391,7 +391,7 @@ struct GPUHistMakerDevice {
   // instances to their final leaf. This information is used later to update the
   // prediction cache
   void FinalisePosition(RegTree const* p_tree, size_t n_leaf, DMatrix* p_fmat, ObjInfo task,
-                        std::vector<RowIndexCache>* p_out_row_indices) {
+                        HostDeviceVector<bst_node_t>* p_out_position) {
     dh::TemporaryArray<RegTree::Node> d_nodes(p_tree->GetNodes().size());
     dh::safe_cuda(cudaMemcpyAsync(d_nodes.data().get(), p_tree->GetNodes().data(),
                                   d_nodes.size() * sizeof(RegTree::Node),
@@ -410,8 +410,6 @@ struct GPUHistMakerDevice {
       dh::CopyToD(categories_segments, &d_categories_segments);
     }
 
-    CHECK(p_out_row_indices->empty());
-    p_out_row_indices->push_back(RowIndexCache{ctx_, p_fmat->Info().num_row_});
     if (row_partitioner->GetRows().size() != p_fmat->Info().num_row_) {
       row_partitioner.reset();  // Release the device memory first before reallocating
       row_partitioner.reset(new RowPartitioner(ctx_->gpu_id, p_fmat->Info().num_row_));
@@ -423,12 +421,12 @@ struct GPUHistMakerDevice {
     if (page->n_rows == p_fmat->Info().num_row_) {
       FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
                              dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
-                             n_leaf, p_out_row_indices);
+                             n_leaf, p_out_position);
     } else {
       for (auto const& batch : p_fmat->GetBatches<EllpackPage>(batch_param)) {
         FinalisePositionInPage(batch.Impl(), p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
                                dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
-                               n_leaf, p_out_row_indices);
+                               n_leaf, p_out_position);
       }
     }
   }
@@ -441,11 +439,11 @@ struct GPUHistMakerDevice {
                               common::Span<RegTree::Segment> categories_segments,
                               ObjInfo task,
                               size_t n_leaf,
-                              std::vector<RowIndexCache>* p_out_row_indices) {
+                              HostDeviceVector<bst_node_t>* p_out_position) {
     auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
     auto d_gpair = this->gpair;
     row_partitioner->FinalisePosition(
-        ctx_, p_tree, n_leaf, task, p_out_row_indices,
+        ctx_, p_tree, n_leaf, task, p_out_position,
         [=] __device__(size_t row_id, int position) {
           // What happens if user prune the tree?
           if (!d_matrix.IsInRange(row_id)) {
@@ -652,7 +650,7 @@ struct GPUHistMakerDevice {
 
   void UpdateTree(HostDeviceVector<GradientPair>* gpair_all, DMatrix* p_fmat, ObjInfo task,
                   RegTree* p_tree, dh::AllReducer* reducer,
-                  std::vector<RowIndexCache>* p_out_row_indices) {
+                  HostDeviceVector<bst_node_t>* p_out_position) {
     auto& tree = *p_tree;
     Driver<GPUExpandEntry> driver(static_cast<TrainParam::TreeGrowPolicy>(param.grow_policy));
 
@@ -714,7 +712,7 @@ struct GPUHistMakerDevice {
     }
 
     monitor.Start("FinalisePosition");
-    this->FinalisePosition(p_tree, num_leaves, p_fmat, task, p_out_row_indices);
+    this->FinalisePosition(p_tree, num_leaves, p_fmat, task, p_out_position);
     monitor.Stop("FinalisePosition");
   }
 };
@@ -737,6 +735,7 @@ class GPUHistMakerSpecialised {
   }
 
   void Update(HostDeviceVector<GradientPair>* gpair, DMatrix* dmat,
+              common::Span<HostDeviceVector<bst_node_t>> out_position,
               const std::vector<RegTree*>& trees) {
     monitor_.Start("Update");
 
@@ -744,17 +743,16 @@ class GPUHistMakerSpecialised {
     float lr = param_.learning_rate;
     param_.learning_rate = lr / trees.size();
 
-    row_set_collection_.clear();
     // build tree
     try {
+      size_t t_idx{0};
       for (xgboost::RegTree* tree : trees) {
-        row_set_collection_.emplace_back();
-        auto& row_indices = row_set_collection_.back();
-        this->UpdateTree(gpair, dmat, tree, &row_indices);
+        this->UpdateTree(gpair, dmat, tree, &out_position[t_idx]);
 
         if (hist_maker_param_.debug_synchronize) {
           this->CheckTreesSynchronized(tree);
         }
+        ++t_idx;
       }
       dh::safe_cuda(cudaGetLastError());
     } catch (const std::exception& e) {
@@ -814,13 +812,13 @@ class GPUHistMakerSpecialised {
   }
 
   void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat, RegTree* p_tree,
-                  std::vector<RowIndexCache>* p_out_row_indices) {
+                  HostDeviceVector<bst_node_t>* p_out_position) {
     monitor_.Start("InitData");
     this->InitData(p_fmat, p_tree);
     monitor_.Stop("InitData");
 
     gpair->SetDevice(ctx_->gpu_id);
-    maker->UpdateTree(gpair, p_fmat, task_, p_tree, &reducer_, p_out_row_indices);
+    maker->UpdateTree(gpair, p_fmat, task_, p_tree, &reducer_, p_out_position);
   }
 
   bool UpdatePredictionCache(const DMatrix *data,
@@ -834,10 +832,6 @@ class GPUHistMakerSpecialised {
     return true;
   }
 
-  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const {
-    return row_set_collection_.at(tree_idx);
-  }
-
   TrainParam param_;   // NOLINT
   MetaInfo* info_{};   // NOLINT
 
@@ -845,7 +839,6 @@ class GPUHistMakerSpecialised {
 
  private:
   bool initialised_ { false };
-  std::vector<std::vector<RowIndexCache>> row_set_collection_;
 
   GPUHistMakerTrainParam hist_maker_param_;
   Context const* ctx_;
@@ -907,11 +900,12 @@ class GPUHistMaker : public TreeUpdater {
   }
 
   void Update(HostDeviceVector<GradientPair>* gpair, DMatrix* dmat,
+              common::Span<HostDeviceVector<bst_node_t>> out_position,
               const std::vector<RegTree*>& trees) override {
     if (hist_maker_param_.single_precision_histogram) {
-      float_maker_->Update(gpair, dmat, trees);
+      float_maker_->Update(gpair, dmat, out_position, trees);
     } else {
-      double_maker_->Update(gpair, dmat, trees);
+      double_maker_->Update(gpair, dmat, out_position, trees);
     }
   }
 
@@ -924,14 +918,6 @@ class GPUHistMaker : public TreeUpdater {
     }
   }
 
-  common::Span<RowIndexCache const> GetRowIndexCache(size_t tree_idx) const override {
-    if (hist_maker_param_.single_precision_histogram) {
-      return float_maker_->GetRowIndexCache(tree_idx);
-    } else {
-      return double_maker_->GetRowIndexCache(tree_idx);
-    }
-  }
-
   char const* Name() const override {
     return "grow_gpu_hist";
   }
diff --git a/tests/cpp/objective/test_regression_obj_gpu.cu b/tests/cpp/objective/test_regression_obj_gpu.cu
index 38f29b8a8800..6115aa12213f 100644
--- a/tests/cpp/objective/test_regression_obj_gpu.cu
+++ b/tests/cpp/objective/test_regression_obj_gpu.cu
@@ -1,6 +1,40 @@
 /*!
- * Copyright 2018 XGBoost contributors
+ * Copyright 2018-2022 by XGBoost contributors
  */
-// Dummy file to keep the CUDA tests.
 
+#include "../../../src/objective/adaptive.cuh"
 #include "test_regression_obj.cc"
+
+namespace xgboost {
+namespace obj {
+void TestFillMissingLeaf() {
+  std::vector<bst_node_t> missing{1, 3};
+  Context ctx;
+
+  HostDeviceVector<bst_node_t> node_idx = {2, 4, 5};
+  HostDeviceVector<size_t> node_ptr = {0, 4, 8, 16};
+  node_idx.SetDevice(0);
+  node_ptr.SetDevice(0);
+
+  detail::FillMissingLeaf(missing, &node_idx, &node_ptr);
+
+  auto const& h_nidx = node_idx.HostVector();
+  auto const& h_nptr = node_ptr.HostVector();
+
+  ASSERT_EQ(h_nidx[0], missing[0]);
+  ASSERT_EQ(h_nidx[2], missing[1]);
+  ASSERT_EQ(h_nidx[1], 2);
+  ASSERT_EQ(h_nidx[3], 4);
+  ASSERT_EQ(h_nidx[4], 5);
+
+  ASSERT_EQ(h_nptr[0], 0);
+  ASSERT_EQ(h_nptr[1], 0);  // empty
+  ASSERT_EQ(h_nptr[2], 4);
+  ASSERT_EQ(h_nptr[3], 4);  // empty
+  ASSERT_EQ(h_nptr[4], 8);
+  ASSERT_EQ(h_nptr[5], 16);
+}
+
+TEST(Adaptive, MissingLeaf) { TestFillMissingLeaf(); }
+}  // namespace obj
+}  // namespace xgboost
diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index 32e6d2cf6a1f..f96796a620ae 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -112,14 +112,14 @@ void TestFinalise() {
   ObjInfo task{ObjInfo::kRegression, false, false};
   RegTree tree;
   tree.ExpandNode(0, 0, 0.f, true, 0., 0., 0., /*loss_chg=*/0.f, 0.f, 0.f, 0.f);
-  std::vector<RowIndexCache> row_index;
+  HostDeviceVector<bst_node_t> position;
   Context ctx;
   ctx.gpu_id = 0;
 
   {
     RowPartitioner rp(0, kNumRows);
     rp.FinalisePosition(
-        &ctx, &tree, tree.GetNumLeaves(), task, &row_index,
+        &ctx, &tree, tree.GetNumLeaves(), task, &position,
         [=] __device__(RowPartitioner::RowIndexT ridx, int position) { return 7; },
         [] XGBOOST_DEVICE(size_t idx) { return false; });
 
@@ -142,68 +142,25 @@ void TestFinalise() {
     }
   }
 
-  row_index.emplace_back(&ctx, kNumRows);
   auto d_hess = dh::ToSpan(hess);
   task.zero_hess = true;
 
   RowPartitioner rp(0, kNumRows);
   rp.FinalisePosition(
-      &ctx, &tree, tree.GetNumLeaves(), task, &row_index,
+      &ctx, &tree, tree.GetNumLeaves(), task, &position,
       [] __device__(RowPartitioner::RowIndexT ridx, bst_node_t position) {
         return ridx % 2 == 0 ? 1 : 2;
       },
       [d_hess] __device__(size_t ridx) { return d_hess[ridx] - 0.f == 0.f; });
 
-  auto const& h_node_ptr = row_index.back().node_ptr.ConstHostVector();
-  ASSERT_EQ(h_node_ptr.size(), 3);
-  ASSERT_EQ(h_node_ptr[0], 4);
-  ASSERT_EQ(h_node_ptr[1], 7);
-  ASSERT_EQ(h_node_ptr[2], kNumRows);
-
-  auto const& h_node_idx = row_index.back().node_idx.ConstHostVector();
-  ASSERT_EQ(h_node_idx.size(), 2);
-  ASSERT_EQ(h_node_idx[0], 1);
-  ASSERT_EQ(h_node_idx[1], 2);
-
-  auto const& h_ridx = row_index.back().row_index.ConstHostVector();
-  std::vector<size_t> sol{0, 3, 6, 9, 2, 4, 8, 1, 5, 7};
-  for (size_t i = 0; i < h_ridx.size(); ++i) {
-    ASSERT_EQ(h_ridx[i], sol[i]);
+  auto const& h_position = position.ConstHostVector();
+  for (auto v : h_position) {
+    std::cout << v << ", ";
   }
+  std::cout << std::endl;
 }
 
-void TestFillMissingLeaf() {
-  std::vector<bst_node_t> missing{1, 3};
-  Context ctx;
-  RowIndexCache row_index(&ctx, 10);
-  row_index.node_idx = {2, 4, 5};
-  row_index.node_ptr = {0, 4, 8, 16};
-  row_index.node_idx.SetDevice(0);
-  row_index.node_ptr.SetDevice(0);
-
-  detail::FillMissingLeaf(missing, &row_index);
-
-  auto const& h_nidx = row_index.node_idx.HostVector();
-  auto const& h_nptr = row_index.node_ptr.HostVector();
-
-  ASSERT_EQ(h_nidx[0], missing[0]);
-  ASSERT_EQ(h_nidx[2], missing[1]);
-  ASSERT_EQ(h_nidx[1], 2);
-  ASSERT_EQ(h_nidx[3], 4);
-  ASSERT_EQ(h_nidx[4], 5);
-
-  ASSERT_EQ(h_nptr[0], 0);
-  ASSERT_EQ(h_nptr[1], 0);  // empty
-  ASSERT_EQ(h_nptr[2], 4);
-  ASSERT_EQ(h_nptr[3], 4);  // empty
-  ASSERT_EQ(h_nptr[4], 8);
-  ASSERT_EQ(h_nptr[5], 16);
-}
-
-TEST(RowPartitioner, Finalise) {
-  TestFillMissingLeaf();
-  TestFinalise();
-}
+TEST(RowPartitioner, Finalise) { TestFinalise(); }
 
 void TestIncorrectRow() {
   RowPartitioner rp(0, 1);
diff --git a/tests/cpp/tree/test_gpu_hist.cu b/tests/cpp/tree/test_gpu_hist.cu
index 2ae08a4e8bb3..6490c1cd0f1c 100644
--- a/tests/cpp/tree/test_gpu_hist.cu
+++ b/tests/cpp/tree/test_gpu_hist.cu
@@ -350,7 +350,8 @@ void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* dmat,
   GenericParameter generic_param(CreateEmptyGenericParam(0));
   hist_maker.Configure(args, &generic_param);
 
-  hist_maker.Update(gpair, dmat, {tree});
+  std::vector<HostDeviceVector<bst_node_t>> position;
+  hist_maker.Update(gpair, dmat, common::Span<HostDeviceVector<bst_node_t>>{position}, {tree});
   auto cache = linalg::VectorView<float>{preds->DeviceSpan(), {preds->Size()}, 0};
   hist_maker.UpdatePredictionCache(dmat, cache);
 }

From d83a97366ac1b4a51edb7139c54ce97372f83769 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 00:16:13 +0800
Subject: [PATCH 109/124] Cleanup.

---
 include/xgboost/tree_model.h                  | 21 -------------------
 src/objective/adaptive.cuh                    | 11 +++++-----
 src/objective/regression_obj.cu               |  3 ++-
 .../cpp/tree/gpu_hist/test_row_partitioner.cu |  9 +++++---
 4 files changed, 14 insertions(+), 30 deletions(-)

diff --git a/include/xgboost/tree_model.h b/include/xgboost/tree_model.h
index 7912da809378..3658efd94ad5 100644
--- a/include/xgboost/tree_model.h
+++ b/include/xgboost/tree_model.h
@@ -735,26 +735,5 @@ inline bool RegTree::FVec::IsMissing(size_t i) const {
 inline bool RegTree::FVec::HasMissing() const {
   return has_missing_;
 }
-
-/**
- * \brief A cache for row partition, each partition is the row index of a tree leaf.
- */
-struct RowIndexCache {
-  HostDeviceVector<size_t> row_index;
-  HostDeviceVector<size_t> node_ptr;
-  HostDeviceVector<bst_node_t> node_idx;
-
-  /**
-   * \param ctx Context
-   * \param n_samples The number of samples for this cache, which equals to the number of
-   *                  samples in a single page from DMatarix.
-   */
-  RowIndexCache(Context const* ctx, size_t n_samples) {
-    if (!ctx->IsCPU()) {
-      row_index.SetDevice(ctx->gpu_id);
-    }
-    row_index.Resize(n_samples);
-  }
-};
 }  // namespace xgboost
 #endif  // XGBOOST_TREE_MODEL_H_
diff --git a/src/objective/adaptive.cuh b/src/objective/adaptive.cuh
index e5a32fc246f6..c551f5899b5b 100644
--- a/src/objective/adaptive.cuh
+++ b/src/objective/adaptive.cuh
@@ -36,8 +36,8 @@ inline void FillMissingLeaf(std::vector<bst_node_t> const& maybe_missing,
 }
 
 inline void EncodeTreeLeaf(Context const* ctx, common::Span<bst_node_t const> position,
-                           HostDeviceVector<size_t>* p_nptr, HostDeviceVector<bst_node_t>* p_nidx,
-                           RegTree const& tree) {
+                           dh::device_vector<size_t>* p_ridx, HostDeviceVector<size_t>* p_nptr,
+                           HostDeviceVector<bst_node_t>* p_nidx, RegTree const& tree) {
   // copy position to buffer
   dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
   size_t n_samples = position.size();
@@ -45,11 +45,12 @@ inline void EncodeTreeLeaf(Context const* ctx, common::Span<bst_node_t const> po
   dh::device_vector<bst_node_t> sorted_position(position.size());
   dh::safe_cuda(cudaMemcpyAsync(sorted_position.data().get(), position.data(),
                                 position.size_bytes(), cudaMemcpyDeviceToDevice));
-  dh::device_vector<size_t> ridx(position.size());
-  dh::Iota(dh::ToSpan(ridx));
+
+  p_ridx->resize(position.size());
+  dh::Iota(dh::ToSpan(*p_ridx));
   // sort row index according to node index
   thrust::stable_sort_by_key(thrust::cuda::par(alloc), sorted_position.begin(),
-                             sorted_position.begin() + n_samples, ridx.begin());
+                             sorted_position.begin() + n_samples, p_ridx->begin());
 
   size_t n_leaf = tree.GetNumLeaves();
   // +1 for subsample, which is set to an unique value in above kernel.
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index a7260965cdec..01d6ad8bbc8f 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -723,7 +723,7 @@ void UpdateTreeLeafDevice(Context const* ctx, common::Span<bst_node_t const> pos
   HostDeviceVector<size_t> nptr;
   HostDeviceVector<bst_node_t> nidx;
 
-  EncodeTreeLeaf(ctx, position, &nptr, &nidx, *p_tree);
+  EncodeTreeLeaf(ctx, position, &ridx, &nptr, &nidx, *p_tree);
 
   HostDeviceVector<float> quantiles;
   predt.SetDevice(ctx->gpu_id);
@@ -902,6 +902,7 @@ class MeanAbsoluteError : public ObjFunction {
       detail::UpdateTreeLeafHost(ctx_, h_position, info, prediction, 0.5, p_tree);
     } else {
 #if defined(XGBOOST_USE_CUDA)
+      position.SetDevice(ctx_->gpu_id);
       auto d_position = position.ConstDeviceSpan();
       detail::UpdateTreeLeafDevice(ctx_, d_position, info, prediction, 0.5, p_tree);
 #else
diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index f96796a620ae..126b51497f8f 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -154,10 +154,13 @@ void TestFinalise() {
       [d_hess] __device__(size_t ridx) { return d_hess[ridx] - 0.f == 0.f; });
 
   auto const& h_position = position.ConstHostVector();
-  for (auto v : h_position) {
-    std::cout << v << ", ";
+  for (size_t ridx = 0; ridx < h_position.size(); ++ridx) {
+    if (ridx % 3 == 0) {
+      ASSERT_LT(h_position[ridx], 0);
+    } else {
+      ASSERT_EQ(h_position[ridx], ridx % 2 == 0 ? 1 : 2);
+    }
   }
-  std::cout << std::endl;
 }
 
 TEST(RowPartitioner, Finalise) { TestFinalise(); }

From a4ce3144a979d589e52d808d8c66ddb347e57421 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 01:30:08 +0800
Subject: [PATCH 110/124] Refactor.

---
 amalgamation/xgboost-all0.cc                  |   1 +
 src/common/partition_builder.h                |   1 +
 src/objective/adaptive.cc                     | 146 +++++++++++++
 src/objective/{adaptive.cuh => adaptive.cu}   | 141 ++++++++-----
 src/objective/adaptive.h                      |  82 ++++++++
 src/objective/regression_obj.cu               | 198 +-----------------
 tests/cpp/objective/test_regression_obj.cc    |  27 ++-
 .../cpp/objective/test_regression_obj_gpu.cu  |  30 +--
 tests/cpp/tree/test_gpu_hist.cu               |   2 +-
 9 files changed, 342 insertions(+), 286 deletions(-)
 create mode 100644 src/objective/adaptive.cc
 rename src/objective/{adaptive.cuh => adaptive.cu} (50%)
 create mode 100644 src/objective/adaptive.h

diff --git a/amalgamation/xgboost-all0.cc b/amalgamation/xgboost-all0.cc
index 45eb5e72593d..c684e6309de8 100644
--- a/amalgamation/xgboost-all0.cc
+++ b/amalgamation/xgboost-all0.cc
@@ -24,6 +24,7 @@
 #include "../src/objective/rank_obj.cc"
 #include "../src/objective/hinge.cc"
 #include "../src/objective/aft_obj.cc"
+#include "../src/objective/adaptive.cc"
 
 // gbms
 #include "../src/gbm/gbm.cc"
diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index 06973979c723..8f2c481dc1b4 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -12,6 +12,7 @@
 #include <algorithm>
 #include <memory>
 #include <utility>
+#include <limits>
 #include <vector>
 
 #include "categorical.h"
diff --git a/src/objective/adaptive.cc b/src/objective/adaptive.cc
new file mode 100644
index 000000000000..f5f854b22d9b
--- /dev/null
+++ b/src/objective/adaptive.cc
@@ -0,0 +1,146 @@
+/*!
+ * Copyright 2022 by XGBoost Contributors
+ */
+#include "adaptive.h"
+
+#include <limits>
+#include <vector>
+
+#include "../common/common.h"
+#include "../common/stats.h"
+#include "../common/threading_utils.h"
+#include "xgboost/tree_model.h"
+
+namespace xgboost {
+namespace obj {
+namespace detail {
+void EncodeTreeLeafHost(RegTree const& tree, std::vector<bst_node_t> const& position,
+                        std::vector<size_t>* p_nptr, std::vector<bst_node_t>* p_nidx,
+                        std::vector<size_t>* p_ridx) {
+  auto& nptr = *p_nptr;
+  auto& nidx = *p_nidx;
+  auto& ridx = *p_ridx;
+  ridx = common::ArgSort<size_t>(position);
+  std::vector<bst_node_t> sorted_pos(position);
+  // permutation
+  for (size_t i = 0; i < position.size(); ++i) {
+    sorted_pos[i] = position[ridx[i]];
+  }
+  // find the first non-sampled row
+  auto begin_pos =
+      std::distance(sorted_pos.cbegin(), std::find_if(sorted_pos.cbegin(), sorted_pos.cend(),
+                                                      [](bst_node_t nidx) { return nidx >= 0; }));
+  CHECK_LE(begin_pos, sorted_pos.size());
+
+  std::vector<bst_node_t> leaf;
+  tree.WalkTree([&](bst_node_t nidx) {
+    if (tree[nidx].IsLeaf()) {
+      leaf.push_back(nidx);
+    }
+    return true;
+  });
+
+  if (begin_pos == sorted_pos.size()) {
+    nidx = leaf;
+    return;
+  }
+
+  auto beg_it = sorted_pos.begin() + begin_pos;
+  common::RunLengthEncode(beg_it, sorted_pos.end(), &nptr);
+  CHECK_GT(nptr.size(), 0);
+  // skip the sampled rows in indptr
+  std::transform(nptr.begin(), nptr.end(), nptr.begin(),
+                 [begin_pos](size_t ptr) { return ptr + begin_pos; });
+
+  size_t n_leaf = nptr.size() - 1;
+  auto n_unique = std::unique(beg_it, sorted_pos.end()) - beg_it;
+  CHECK_EQ(n_unique, n_leaf);
+  nidx.resize(n_leaf);
+  std::copy(beg_it, beg_it + n_unique, nidx.begin());
+
+  if (n_leaf != leaf.size()) {
+    FillMissingLeaf(leaf, &nidx, &nptr);
+    n_leaf = leaf.size();
+  }
+}
+
+void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& position,
+                        MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
+                        RegTree* p_tree) {
+  auto& tree = *p_tree;
+  CHECK(!position.empty());
+
+  std::vector<bst_node_t> nidx;
+  std::vector<size_t> nptr;
+  std::vector<size_t> ridx;
+  EncodeTreeLeafHost(*p_tree, position, &nptr, &nidx, &ridx);
+  size_t n_leaf = nidx.size();
+  if (nptr.empty()) {
+    std::vector<float> quantiles;
+    UpdateLeafValues(&quantiles, nidx, p_tree);
+  }
+
+  std::vector<float> quantiles(n_leaf, 0);
+  std::vector<int32_t> n_valids(n_leaf, 0);
+
+  {
+    std::vector<float> results(nidx.size());
+    auto const& h_node_idx = nidx;
+    auto const& h_node_ptr = nptr;
+    CHECK_LE(h_node_ptr.back(), info.num_row_);
+    // loop over each leaf
+    common::ParallelFor(results.size(), ctx->Threads(), [&](size_t k) {
+      auto nidx = h_node_idx[k];
+      CHECK(tree[nidx].IsLeaf());
+      CHECK_LT(k + 1, h_node_ptr.size());
+      size_t n = h_node_ptr[k + 1] - h_node_ptr[k];
+      auto h_row_set = common::Span<size_t const>{ridx}.subspan(h_node_ptr[k], n);
+      // multi-target not yet supported.
+      auto h_labels = info.labels.HostView().Slice(linalg::All(), 0);
+      auto const& h_predt = predt.ConstHostVector();
+      auto h_weights = linalg::MakeVec(&info.weights_);
+
+      auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
+        auto row_idx = h_row_set[i];
+        return h_labels(row_idx) - h_predt[row_idx];
+      });
+      auto w_it = common::MakeIndexTransformIter([&](size_t i) -> float {
+        auto row_idx = h_row_set[i];
+        return h_weights(row_idx);
+      });
+
+      float q{0};
+      if (info.weights_.Empty()) {
+        q = common::Quantile(alpha, iter, iter + h_row_set.size());
+      } else {
+        q = common::WeightedQuantile(alpha, iter, iter + h_row_set.size(), w_it);
+      }
+      if (std::isnan(q)) {
+        CHECK(h_row_set.empty());
+      }
+      results.at(k) = q;
+    });
+
+    // sum result from each external memory partition to quantiles
+    for (size_t i = 0; i < results.size(); ++i) {
+      if (!std::isnan(results[i])) {
+        quantiles[i] += results[i];
+        n_valids[i]++;
+      }
+    }
+  }
+
+  for (size_t i = 0; i < quantiles.size(); ++i) {
+    if (n_valids[i] > 0) {
+      quantiles[i] /= n_valids[i];
+    } else {
+      // mark that no page has valid sample in the i^th leaf
+      quantiles[i] = std::numeric_limits<float>::quiet_NaN();
+    }
+  }
+
+  UpdateLeafValues(&quantiles, nidx, p_tree);
+}
+}  // namespace detail
+}  // namespace obj
+}  // namespace xgboost
diff --git a/src/objective/adaptive.cuh b/src/objective/adaptive.cu
similarity index 50%
rename from src/objective/adaptive.cuh
rename to src/objective/adaptive.cu
index c551f5899b5b..42d239acd977 100644
--- a/src/objective/adaptive.cuh
+++ b/src/objective/adaptive.cu
@@ -1,43 +1,20 @@
 /*!
  * Copyright 2022 by XGBoost Contributors
  */
-#pragma once
 #include <thrust/sort.h>
 
 #include <cub/cub.cuh>
 
 #include "../common/device_helpers.cuh"
-#include "xgboost/generic_parameters.h"
-#include "xgboost/host_device_vector.h"
-#include "xgboost/tree_model.h"
+#include "../common/stats.cuh"
+#include "adaptive.h"
 
 namespace xgboost {
 namespace obj {
 namespace detail {
-inline void FillMissingLeaf(std::vector<bst_node_t> const& maybe_missing,
-                            HostDeviceVector<bst_node_t>* p_nidx,
-                            HostDeviceVector<size_t>* p_nptr) {
-  auto& h_node_idx = p_nidx->HostVector();
-  auto& h_node_ptr = p_nptr->HostVector();
-
-  for (auto leaf : maybe_missing) {
-    if (std::binary_search(h_node_idx.cbegin(), h_node_idx.cend(), leaf)) {
-      continue;
-    }
-    auto it = std::upper_bound(h_node_idx.cbegin(), h_node_idx.cend(), leaf);
-    auto pos = it - h_node_idx.cbegin();
-    h_node_idx.insert(h_node_idx.cbegin() + pos, leaf);
-    h_node_ptr.insert(h_node_ptr.cbegin() + pos, h_node_ptr[pos]);
-  }
-
-  // push to device.
-  p_nidx->ConstDevicePointer();
-  p_nptr->ConstDevicePointer();
-}
-
-inline void EncodeTreeLeaf(Context const* ctx, common::Span<bst_node_t const> position,
-                           dh::device_vector<size_t>* p_ridx, HostDeviceVector<size_t>* p_nptr,
-                           HostDeviceVector<bst_node_t>* p_nidx, RegTree const& tree) {
+void EncodeTreeLeafDevice(Context const* ctx, common::Span<bst_node_t const> position,
+                          dh::device_vector<size_t>* p_ridx, HostDeviceVector<size_t>* p_nptr,
+                          HostDeviceVector<bst_node_t>* p_nidx, RegTree const& tree) {
   // copy position to buffer
   dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
   size_t n_samples = position.size();
@@ -51,10 +28,24 @@ inline void EncodeTreeLeaf(Context const* ctx, common::Span<bst_node_t const> po
   // sort row index according to node index
   thrust::stable_sort_by_key(thrust::cuda::par(alloc), sorted_position.begin(),
                              sorted_position.begin() + n_samples, p_ridx->begin());
+  dh::XGBCachingDeviceAllocator<char> caching;
+  auto beg_pos =
+      thrust::find_if(thrust::cuda::par(caching), sorted_position.cbegin(), sorted_position.cend(),
+                      [] XGBOOST_DEVICE(bst_node_t nidx) { return nidx >= 0; }) -
+      sorted_position.cbegin();
+  if (beg_pos == sorted_position.size()) {
+    auto& leaf = p_nidx->HostVector();
+    tree.WalkTree([&](bst_node_t nidx) {
+      if (tree[nidx].IsLeaf()) {
+        leaf.push_back(nidx);
+      }
+      return true;
+    });
+    return;
+  }
 
   size_t n_leaf = tree.GetNumLeaves();
-  // +1 for subsample, which is set to an unique value in above kernel.
-  size_t max_n_unique = n_leaf + 1;
+  size_t max_n_unique = n_leaf;
 
   dh::caching_device_vector<size_t> counts_out(max_n_unique + 1, 0);
   auto d_counts_out = dh::ToSpan(counts_out).subspan(0, max_n_unique);
@@ -63,15 +54,15 @@ inline void EncodeTreeLeaf(Context const* ctx, common::Span<bst_node_t const> po
   auto d_unique_out = dh::ToSpan(unique_out);
 
   size_t nbytes;
-  cub::DeviceRunLengthEncode::Encode(nullptr, nbytes, sorted_position.begin(),
-                                     unique_out.data().get(), counts_out.data().get(),
-                                     d_num_runs_out.data(), n_samples);
+  auto begin_it = sorted_position.begin() + beg_pos;
+  cub::DeviceRunLengthEncode::Encode(nullptr, nbytes, begin_it, unique_out.data().get(),
+                                     counts_out.data().get(), d_num_runs_out.data(),
+                                     n_samples - beg_pos);
   dh::TemporaryArray<char> temp(nbytes);
-  cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, sorted_position.begin(),
-                                     unique_out.data().get(), counts_out.data().get(),
-                                     d_num_runs_out.data(), n_samples);
+  cub::DeviceRunLengthEncode::Encode(temp.data().get(), nbytes, begin_it, unique_out.data().get(),
+                                     counts_out.data().get(), d_num_runs_out.data(),
+                                     n_samples - beg_pos);
 
-  dh::XGBCachingDeviceAllocator<char> caching;
   dh::PinnedMemory pinned_pool;
   auto pinned = pinned_pool.GetSpan<char>(sizeof(size_t) + sizeof(bst_node_t));
   dh::CUDAStream copy_stream;
@@ -104,35 +95,20 @@ inline void EncodeTreeLeaf(Context const* ctx, common::Span<bst_node_t const> po
       // missing data, which can be caused by sparse input and distributed training.
       return;
     }
-    if (d_unique_out[0] < 0) {
-      // shift 1 to the left
-      // some rows are ignored due to sampling, `kIgnoredTreePosition` is -1 so it's the
-      // smallest value and is sorted to the left.
-      // d_unique_out.size() == n_leaf + 1.
-      d_node_idx[i] = d_unique_out[i + 1];
-      d_node_ptr[i + 1] = d_counts_out[i + 1];
-      if (i == 0) {
-        d_node_ptr[0] = d_counts_out[0];
-      }
-    } else {
-      d_node_idx[i] = d_unique_out[i];
-      d_node_ptr[i + 1] = d_counts_out[i];
-      if (i == 0) {
-        d_node_ptr[0] = 0;
-      }
+    d_node_idx[i] = d_unique_out[i];
+    d_node_ptr[i + 1] = d_counts_out[i];
+    if (i == 0) {
+      d_node_ptr[0] = beg_pos;
     }
   });
   thrust::inclusive_scan(thrust::cuda::par(caching), dh::tbegin(d_node_ptr), dh::tend(d_node_ptr),
                          dh::tbegin(d_node_ptr));
   copy_stream.View().Sync();
-  if (*h_first_unique < 0) {
-    *h_num_runs -= 1;  // sampled.
-  }
   CHECK_GT(*h_num_runs, 0);
   CHECK_LE(*h_num_runs, n_leaf);
 
   if (*h_num_runs < n_leaf) {
-    // shrink to omit the `kIgnoredTreePosition`.
+    // shrink to omit the sampled nodes.
     nptr.Resize(*h_num_runs + 1);
     nidx.Resize(*h_num_runs);
 
@@ -147,11 +123,60 @@ inline void EncodeTreeLeaf(Context const* ctx, common::Span<bst_node_t const> po
     // Fill all the leaves that don't have any sample. This is hacky and inefficient. An
     // alternative is to leave the objective to handle missing leaf, which is more messy
     // as we need to take other distributed workers into account.
-    FillMissingLeaf(leaves, &nidx, &nptr);
+    auto& h_nidx = nidx.HostVector();
+    auto& h_nptr = nptr.HostVector();
+    FillMissingLeaf(leaves, &h_nidx, &h_nptr);
+    nidx.DevicePointer();
+    nptr.DevicePointer();
   }
   CHECK_EQ(nidx.Size(), n_leaf);
   CHECK_EQ(nptr.Size(), n_leaf + 1);
 }
+
+void UpdateTreeLeafDevice(Context const* ctx, common::Span<bst_node_t const> position,
+                          MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
+                          RegTree* p_tree) {
+  dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
+  dh::device_vector<size_t> ridx;
+  HostDeviceVector<size_t> nptr;
+  HostDeviceVector<bst_node_t> nidx;
+
+  EncodeTreeLeafDevice(ctx, position, &ridx, &nptr, &nidx, *p_tree);
+
+  if (nptr.Empty()) {
+    std::vector<float> quantiles;
+    UpdateLeafValues(&quantiles, nidx.ConstHostVector(), p_tree);
+  }
+
+  HostDeviceVector<float> quantiles;
+  predt.SetDevice(ctx->gpu_id);
+  auto d_predt = predt.ConstDeviceSpan();
+  auto d_labels = info.labels.View(ctx->gpu_id);
+
+  auto d_row_index = dh::ToSpan(ridx);
+  auto seg_beg = nptr.DevicePointer();
+  auto seg_end = seg_beg + nptr.Size();
+  auto val_beg = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
+                                                  [=] XGBOOST_DEVICE(size_t i) {
+                                                    auto predt = d_predt[d_row_index[i]];
+                                                    auto y = d_labels(d_row_index[i]);
+                                                    return y - predt;
+                                                  });
+  auto val_end = val_beg + d_labels.Size();
+  CHECK_EQ(nidx.Size() + 1, nptr.Size());
+  if (info.weights_.Empty()) {
+    common::SegmentedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, &quantiles);
+  } else {
+    info.weights_.SetDevice(ctx->gpu_id);
+    auto d_weights = info.weights_.ConstDeviceSpan();
+    CHECK_EQ(d_weights.size(), d_row_index.size());
+    auto w_it = thrust::make_permutation_iterator(dh::tcbegin(d_weights), dh::tcbegin(d_row_index));
+    common::SegmentedWeightedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, w_it,
+                                      w_it + d_weights.size(), &quantiles);
+  }
+
+  UpdateLeafValues(&quantiles.HostVector(), nidx.ConstHostVector(), p_tree);
+}
 }  // namespace detail
 }  // namespace obj
 }  // namespace xgboost
diff --git a/src/objective/adaptive.h b/src/objective/adaptive.h
new file mode 100644
index 000000000000..89d29a545cdc
--- /dev/null
+++ b/src/objective/adaptive.h
@@ -0,0 +1,82 @@
+/*!
+ * Copyright 2022 by XGBoost Contributors
+ */
+#pragma once
+
+#include <algorithm>
+#include <vector>
+
+#include "rabit/rabit.h"
+#include "xgboost/generic_parameters.h"
+#include "xgboost/host_device_vector.h"
+#include "xgboost/tree_model.h"
+
+namespace xgboost {
+namespace obj {
+namespace detail {
+inline void FillMissingLeaf(std::vector<bst_node_t> const& maybe_missing,
+                            std::vector<bst_node_t>* p_nidx, std::vector<size_t>* p_nptr) {
+  auto& h_node_idx = *p_nidx;
+  auto& h_node_ptr = *p_nptr;
+
+  for (auto leaf : maybe_missing) {
+    if (std::binary_search(h_node_idx.cbegin(), h_node_idx.cend(), leaf)) {
+      continue;
+    }
+    auto it = std::upper_bound(h_node_idx.cbegin(), h_node_idx.cend(), leaf);
+    auto pos = it - h_node_idx.cbegin();
+    h_node_idx.insert(h_node_idx.cbegin() + pos, leaf);
+    h_node_ptr.insert(h_node_ptr.cbegin() + pos, h_node_ptr[pos]);
+  }
+}
+
+inline void UpdateLeafValues(std::vector<float>* p_quantiles, std::vector<bst_node_t> const nidx,
+                             RegTree* p_tree) {
+  auto& tree = *p_tree;
+  auto& quantiles = *p_quantiles;
+  auto const& h_node_idx = nidx;
+
+  size_t n_leaf{h_node_idx.size()};
+  rabit::Allreduce<rabit::op::Max>(&n_leaf, 1);
+  CHECK(quantiles.empty() || quantiles.size() == n_leaf);
+  if (quantiles.empty()) {
+    quantiles.resize(n_leaf);
+  }
+
+  // number of workers that have valid quantiles
+  std::vector<int32_t> n_valids(quantiles.size());
+  std::transform(quantiles.cbegin(), quantiles.cend(), n_valids.begin(),
+                 [](float q) { return static_cast<int32_t>(!std::isnan(q)); });
+  rabit::Allreduce<rabit::op::Sum>(n_valids.data(), n_valids.size());
+  // convert to 0 for all reduce
+  std::replace_if(
+      quantiles.begin(), quantiles.end(), [](float q) { return std::isnan(q); }, 0.f);
+  // use the mean value
+  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
+  for (size_t i = 0; i < n_leaf; ++i) {
+    if (n_valids[i] > 0) {
+      quantiles[i] /= static_cast<float>(n_valids[i]);
+    } else {
+      // Use original leaf value if no worker can provide the quantile.
+      quantiles[i] = tree[h_node_idx[i]].LeafValue();
+    }
+  }
+
+  for (size_t i = 0; i < nidx.size(); ++i) {
+    auto nidx = h_node_idx[i];
+    auto q = quantiles[i];
+    CHECK(tree[nidx].IsLeaf());
+    tree[nidx].SetLeaf(q);
+  }
+}
+
+void UpdateTreeLeafDevice(Context const* ctx, common::Span<bst_node_t const> position,
+                          MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
+                          RegTree* p_tree);
+
+void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& position,
+                        MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
+                        RegTree* p_tree);
+}  // namespace detail
+}  // namespace obj
+}  // namespace xgboost
diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 01d6ad8bbc8f..7f81322f2929 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -4,7 +4,6 @@
  * \brief Definition of single-value regression and classification objectives.
  * \author Tianqi Chen, Kailong Chen
  */
-
 #include <dmlc/omp.h>
 #include <rabit/rabit.h>
 #include <xgboost/logging.h>
@@ -18,10 +17,10 @@
 #include "../common/common.h"
 #include "../common/linalg_op.h"
 #include "../common/pseudo_huber.h"
-#include "../common/stats.h"
 #include "../common/threading_utils.h"
 #include "../common/transform.h"
 #include "./regression_loss.h"
+#include "adaptive.h"
 #include "xgboost/base.h"
 #include "xgboost/data.h"
 #include "xgboost/generic_parameters.h"
@@ -34,8 +33,6 @@
 #if defined(XGBOOST_USE_CUDA)
 #include "../common/device_helpers.cuh"
 #include "../common/linalg_op.cuh"
-#include "../common/stats.cuh"
-#include "adaptive.cuh"
 #endif  // defined(XGBOOST_USE_CUDA)
 
 namespace xgboost {
@@ -671,199 +668,6 @@ XGBOOST_REGISTER_OBJECTIVE(TweedieRegression, "reg:tweedie")
 .describe("Tweedie regression for insurance data.")
 .set_body([]() { return new TweedieRegression(); });
 
-namespace detail {
-void UpdateLeafValues(std::vector<float>* p_quantiles, std::vector<bst_node_t> const nidx,
-                      RegTree* p_tree) {
-  auto& tree = *p_tree;
-  auto& quantiles = *p_quantiles;
-  auto const& h_node_idx = nidx;
-
-  size_t n_leaf{h_node_idx.size()};
-  rabit::Allreduce<rabit::op::Max>(&n_leaf, 1);
-  CHECK(quantiles.empty() || quantiles.size() == n_leaf);
-  if (quantiles.empty()) {
-    quantiles.resize(n_leaf);
-  }
-
-  // number of workers that have valid quantiles
-  std::vector<int32_t> n_valids(quantiles.size());
-  std::transform(quantiles.cbegin(), quantiles.cend(), n_valids.begin(),
-                 [](float q) { return static_cast<int32_t>(!std::isnan(q)); });
-  rabit::Allreduce<rabit::op::Sum>(n_valids.data(), n_valids.size());
-  // convert to 0 for all reduce
-  std::replace_if(
-      quantiles.begin(), quantiles.end(), [](float q) { return std::isnan(q); }, 0.f);
-  // use the mean value
-  rabit::Allreduce<rabit::op::Sum>(quantiles.data(), quantiles.size());
-  for (size_t i = 0; i < n_leaf; ++i) {
-    if (n_valids[i] > 0) {
-      quantiles[i] /= static_cast<float>(n_valids[i]);
-    } else {
-      // Use original leaf value if no worker can provide the quantile.
-      quantiles[i] = tree[h_node_idx[i]].LeafValue();
-    }
-  }
-
-  for (size_t i = 0; i < nidx.size(); ++i) {
-    auto nidx = h_node_idx[i];
-    auto q = quantiles[i];
-    CHECK(tree[nidx].IsLeaf());
-    tree[nidx].SetLeaf(q);
-  }
-}
-
-#if defined(XGBOOST_USE_CUDA)
-
-
-void UpdateTreeLeafDevice(Context const* ctx, common::Span<bst_node_t const> position,
-                          MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
-                          RegTree* p_tree) {
-  dh::safe_cuda(cudaSetDevice(ctx->gpu_id));
-  dh::device_vector<size_t> ridx;
-  HostDeviceVector<size_t> nptr;
-  HostDeviceVector<bst_node_t> nidx;
-
-  EncodeTreeLeaf(ctx, position, &ridx, &nptr, &nidx, *p_tree);
-
-  HostDeviceVector<float> quantiles;
-  predt.SetDevice(ctx->gpu_id);
-  auto d_predt = predt.ConstDeviceSpan();
-  auto d_labels = info.labels.View(ctx->gpu_id);
-
-
-  auto d_row_index = dh::ToSpan(ridx);
-  auto seg_beg = nptr.DevicePointer();
-  auto seg_end = seg_beg + nptr.Size();
-  auto val_beg = dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
-                                                  [=] XGBOOST_DEVICE(size_t i) {
-                                                    auto predt = d_predt[d_row_index[i]];
-                                                    auto y = d_labels(d_row_index[i]);
-                                                    return y - predt;
-                                                  });
-  auto val_end = val_beg + d_labels.Size();
-  CHECK_EQ(nidx.Size() + 1, nptr.Size());
-  if (info.weights_.Empty()) {
-    common::SegmentedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, &quantiles);
-  } else {
-    info.weights_.SetDevice(ctx->gpu_id);
-    auto d_weights = info.weights_.ConstDeviceSpan();
-    CHECK_EQ(d_weights.size(), d_row_index.size());
-    auto w_it = thrust::make_permutation_iterator(dh::tcbegin(d_weights), dh::tcbegin(d_row_index));
-    common::SegmentedWeightedQuantile(ctx, alpha, seg_beg, seg_end, val_beg, val_end, w_it,
-                                      w_it + d_weights.size(), &quantiles);
-  }
-
-  UpdateLeafValues(&quantiles.HostVector(), nidx.ConstHostVector(), p_tree);
-}
-#endif  // defined(XGBOOST_USE_CUDA)
-
-void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& position,
-                        MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
-                        RegTree* p_tree) {
-  auto& tree = *p_tree;
-  CHECK(!position.empty());
-
-  auto ridx = common::ArgSort<size_t>(position);
-  std::vector<bst_node_t> sorted_pos(position);
-  // permutation
-  for (size_t i = 0; i < position.size(); ++i) {
-    sorted_pos[i] = position[ridx[i]];
-  }
-  // find the first non-sampled row
-  auto begin_pos =
-      std::distance(sorted_pos.cbegin(), std::find_if(sorted_pos.cbegin(), sorted_pos.cend(),
-                                                      [](bst_node_t nidx) { return nidx >= 0; }));
-  CHECK_LE(begin_pos, sorted_pos.size());
-  if (begin_pos == sorted_pos.size()) {
-    std::vector<bst_node_t> leaf;
-    tree.WalkTree([&](bst_node_t nidx) {
-      if (tree[nidx].IsLeaf()) {
-        leaf.push_back(nidx);
-      }
-      return true;
-    });
-    std::vector<float> quantiles;
-    UpdateLeafValues(&quantiles, leaf, p_tree);
-    return;
-  }
-
-  std::vector<size_t> segments;
-  auto beg_it = sorted_pos.begin() + begin_pos;
-  common::RunLengthEncode(beg_it, sorted_pos.end(), &segments);
-  CHECK_GT(segments.size(), 0);
-  // skip the sampled rows in indptr
-  std::transform(segments.begin(), segments.end(), segments.begin(),
-                 [begin_pos](size_t ptr) { return ptr + begin_pos; });
-
-  size_t n_leaf = segments.size() - 1;
-  auto n_unique = std::unique(beg_it, sorted_pos.end()) - beg_it;
-  CHECK_EQ(n_unique, n_leaf);
-  std::vector<bst_node_t> nidx(n_leaf);
-  std::copy(beg_it, beg_it + n_unique, nidx.begin());
-
-  std::vector<float> quantiles(n_leaf, 0);
-  std::vector<int32_t> n_valids(n_leaf, 0);
-
-  {
-    std::vector<float> results(nidx.size());
-    auto const& h_node_idx = nidx;
-    auto const& h_node_ptr = segments;
-    CHECK_LE(h_node_ptr.back(), info.num_row_);
-    // loop over each leaf
-    common::ParallelFor(results.size(), ctx->Threads(), [&](size_t k) {
-      auto nidx = h_node_idx[k];
-      CHECK(tree[nidx].IsLeaf());
-      CHECK_LT(k + 1, h_node_ptr.size());
-      size_t n = h_node_ptr[k + 1] - h_node_ptr[k];
-      auto h_row_set = common::Span<size_t const>{ridx}.subspan(h_node_ptr[k], n);
-      // multi-target not yet supported.
-      auto h_labels = info.labels.HostView().Slice(linalg::All(), 0);
-      auto const& h_predt = predt.ConstHostVector();
-      auto h_weights = linalg::MakeVec(&info.weights_);
-
-      auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
-        auto row_idx = h_row_set[i];
-        return h_labels(row_idx) - h_predt[row_idx];
-      });
-      auto w_it = common::MakeIndexTransformIter([&](size_t i) -> float {
-        auto row_idx = h_row_set[i];
-        return h_weights(row_idx);
-      });
-
-      float q{0};
-      if (info.weights_.Empty()) {
-        q = common::Quantile(alpha, iter, iter + h_row_set.size());
-      } else {
-        q = common::WeightedQuantile(alpha, iter, iter + h_row_set.size(), w_it);
-      }
-      if (std::isnan(q)) {
-        CHECK(h_row_set.empty());
-      }
-      results.at(k) = q;
-    });
-
-    // sum result from each external memory partition to quantiles
-    for (size_t i = 0; i < results.size(); ++i) {
-      if (!std::isnan(results[i])) {
-        quantiles[i] += results[i];
-        n_valids[i]++;
-      }
-    }
-  }
-
-  for (size_t i = 0; i < quantiles.size(); ++i) {
-    if (n_valids[i] > 0) {
-      quantiles[i] /= n_valids[i];
-    } else {
-      // mark that no page has valid sample in the i^th leaf
-      quantiles[i] = std::numeric_limits<float>::quiet_NaN();
-    }
-  }
-
-  UpdateLeafValues(&quantiles, nidx, p_tree);
-}
-}  // namespace detail
-
 class MeanAbsoluteError : public ObjFunction {
  public:
   void Configure(Args const&) override {}
diff --git a/tests/cpp/objective/test_regression_obj.cc b/tests/cpp/objective/test_regression_obj.cc
index 964bb92e41c9..a26f69476152 100644
--- a/tests/cpp/objective/test_regression_obj.cc
+++ b/tests/cpp/objective/test_regression_obj.cc
@@ -2,10 +2,13 @@
  * Copyright 2017-2022 XGBoost contributors
  */
 #include <gtest/gtest.h>
-#include <xgboost/objective.h>
 #include <xgboost/generic_parameters.h>
 #include <xgboost/json.h>
+#include <xgboost/objective.h>
+
+#include "../../../src/objective/adaptive.h"
 #include "../helpers.h"
+
 namespace xgboost {
 
 TEST(Objective, DeclareUnifiedTest(LinearRegressionGPair)) {
@@ -465,4 +468,26 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteErrorLeaf)) {
   ASSERT_EQ(tree[5].LeafValue(), -10);
   ASSERT_EQ(tree[6].LeafValue(), -14);
 }
+
+TEST(Adaptive, DeclareUnifiedTest(MissingLeaf)) {
+  std::vector<bst_node_t> missing{1, 3};
+
+  std::vector<bst_node_t> h_nidx = {2, 4, 5};
+  std::vector<size_t> h_nptr = {0, 4, 8, 16};
+
+  obj::detail::FillMissingLeaf(missing, &h_nidx, &h_nptr);
+
+  ASSERT_EQ(h_nidx[0], missing[0]);
+  ASSERT_EQ(h_nidx[2], missing[1]);
+  ASSERT_EQ(h_nidx[1], 2);
+  ASSERT_EQ(h_nidx[3], 4);
+  ASSERT_EQ(h_nidx[4], 5);
+
+  ASSERT_EQ(h_nptr[0], 0);
+  ASSERT_EQ(h_nptr[1], 0);  // empty
+  ASSERT_EQ(h_nptr[2], 4);
+  ASSERT_EQ(h_nptr[3], 4);  // empty
+  ASSERT_EQ(h_nptr[4], 8);
+  ASSERT_EQ(h_nptr[5], 16);
+}
 }  // namespace xgboost
diff --git a/tests/cpp/objective/test_regression_obj_gpu.cu b/tests/cpp/objective/test_regression_obj_gpu.cu
index 6115aa12213f..b6836207b5ef 100644
--- a/tests/cpp/objective/test_regression_obj_gpu.cu
+++ b/tests/cpp/objective/test_regression_obj_gpu.cu
@@ -2,39 +2,11 @@
  * Copyright 2018-2022 by XGBoost contributors
  */
 
-#include "../../../src/objective/adaptive.cuh"
+
 #include "test_regression_obj.cc"
 
 namespace xgboost {
 namespace obj {
-void TestFillMissingLeaf() {
-  std::vector<bst_node_t> missing{1, 3};
-  Context ctx;
-
-  HostDeviceVector<bst_node_t> node_idx = {2, 4, 5};
-  HostDeviceVector<size_t> node_ptr = {0, 4, 8, 16};
-  node_idx.SetDevice(0);
-  node_ptr.SetDevice(0);
-
-  detail::FillMissingLeaf(missing, &node_idx, &node_ptr);
-
-  auto const& h_nidx = node_idx.HostVector();
-  auto const& h_nptr = node_ptr.HostVector();
-
-  ASSERT_EQ(h_nidx[0], missing[0]);
-  ASSERT_EQ(h_nidx[2], missing[1]);
-  ASSERT_EQ(h_nidx[1], 2);
-  ASSERT_EQ(h_nidx[3], 4);
-  ASSERT_EQ(h_nidx[4], 5);
-
-  ASSERT_EQ(h_nptr[0], 0);
-  ASSERT_EQ(h_nptr[1], 0);  // empty
-  ASSERT_EQ(h_nptr[2], 4);
-  ASSERT_EQ(h_nptr[3], 4);  // empty
-  ASSERT_EQ(h_nptr[4], 8);
-  ASSERT_EQ(h_nptr[5], 16);
-}
 
-TEST(Adaptive, MissingLeaf) { TestFillMissingLeaf(); }
 }  // namespace obj
 }  // namespace xgboost
diff --git a/tests/cpp/tree/test_gpu_hist.cu b/tests/cpp/tree/test_gpu_hist.cu
index 6490c1cd0f1c..3c93c283917a 100644
--- a/tests/cpp/tree/test_gpu_hist.cu
+++ b/tests/cpp/tree/test_gpu_hist.cu
@@ -350,7 +350,7 @@ void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* dmat,
   GenericParameter generic_param(CreateEmptyGenericParam(0));
   hist_maker.Configure(args, &generic_param);
 
-  std::vector<HostDeviceVector<bst_node_t>> position;
+  std::vector<HostDeviceVector<bst_node_t>> position(1);
   hist_maker.Update(gpair, dmat, common::Span<HostDeviceVector<bst_node_t>>{position}, {tree});
   auto cache = linalg::VectorView<float>{preds->DeviceSpan(), {preds->Size()}, 0};
   hist_maker.UpdatePredictionCache(dmat, cache);

From 91c24bbef2741130e65fad5bcc7914648cabaf1b Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 01:44:56 +0800
Subject: [PATCH 111/124] empty.

---
 src/objective/adaptive.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/objective/adaptive.cc b/src/objective/adaptive.cc
index f5f854b22d9b..c43b6a1677d5 100644
--- a/src/objective/adaptive.cc
+++ b/src/objective/adaptive.cc
@@ -68,7 +68,6 @@ void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& posit
                         MetaInfo const& info, HostDeviceVector<float> const& predt, float alpha,
                         RegTree* p_tree) {
   auto& tree = *p_tree;
-  CHECK(!position.empty());
 
   std::vector<bst_node_t> nidx;
   std::vector<size_t> nptr;
@@ -78,8 +77,10 @@ void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& posit
   if (nptr.empty()) {
     std::vector<float> quantiles;
     UpdateLeafValues(&quantiles, nidx, p_tree);
+    return;
   }
 
+  CHECK(!position.empty());
   std::vector<float> quantiles(n_leaf, 0);
   std::vector<int32_t> n_valids(n_leaf, 0);
 

From f86ca051ed69c8cded79dcfb2456edc2b830eb3a Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 01:56:49 +0800
Subject: [PATCH 112/124] Cleanup.

---
 cmake/Utils.cmake              | 6 ------
 include/xgboost/tree_updater.h | 2 ++
 src/common/partition_builder.h | 7 ++++---
 src/gbm/gbtree.h               | 3 ++-
 4 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/cmake/Utils.cmake b/cmake/Utils.cmake
index 687d8e2ebc58..963c494ccf26 100644
--- a/cmake/Utils.cmake
+++ b/cmake/Utils.cmake
@@ -152,12 +152,6 @@ function(xgboost_set_cuda_flags target)
       $<$<COMPILE_LANGUAGE:CUDA>:-lineinfo>)
   endif (USE_DEVICE_DEBUG)
 
-  if (FORCE_COLORED_OUTPUT AND (CMAKE_GENERATOR STREQUAL "Ninja") AND
-      ((CMAKE_CXX_COMPILER_ID STREQUAL "GNU") OR
-        (CMAKE_CXX_COMPILER_ID STREQUAL "Clang")))
-    target_compile_options(${target} PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=-fdiagnostics-color=always>)
-  endif()
-
   if (USE_NVTX)
     enable_nvtx(${target})
   endif (USE_NVTX)
diff --git a/include/xgboost/tree_updater.h b/include/xgboost/tree_updater.h
index a5352aa0a4cf..93cb14afa279 100644
--- a/include/xgboost/tree_updater.h
+++ b/include/xgboost/tree_updater.h
@@ -53,6 +53,8 @@ class TreeUpdater : public Configurable {
    * \brief perform update to the tree models
    * \param gpair the gradient pair statistics of the data
    * \param data The data matrix passed to the updater.
+   * \param out_position The leaf index for each row.  The index is negated if that row is
+   *                     removed during sampling. So the 3th node is ~3.
    * \param out_trees references the trees to be updated, updater will change the content of trees
    *   note: all the trees in the vector are updated, with the same statistics,
    *         but maybe different random seeds, usually one tree is passed in at a time,
diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index 8f2c481dc1b4..5e05a18d0ab4 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -289,9 +289,10 @@ class PartitionBuilder {
     h_pos.resize(row_set.Data()->size(), std::numeric_limits<bst_node_t>::max());
 
     auto p_begin = row_set.Data()->data();
-    for (auto node : row_set) {
+    ParallelFor(row_set.Size(), ctx->Threads(), [&](size_t i) {
+      auto const& node = row_set[i];
       if (node.node_id < 0 || !tree[node.node_id].IsLeaf()) {
-        continue;
+        return;
       }
       if (node.begin) {  // guard for empty node.
         size_t ptr_offset = node.end - p_begin;
@@ -300,7 +301,7 @@ class PartitionBuilder {
           h_pos[*idx] = sampledp(*idx) ? ~node.node_id : node.node_id;
         }
       }
-    }
+    });
   }
 
  protected:
diff --git a/src/gbm/gbtree.h b/src/gbm/gbtree.h
index 885f9cd33db0..020b7d0cb9c0 100644
--- a/src/gbm/gbtree.h
+++ b/src/gbm/gbtree.h
@@ -441,7 +441,8 @@ class GBTree : public GradientBooster {
   Args cfg_;
   // the updaters that can be applied to each of tree
   std::vector<std::unique_ptr<TreeUpdater>> updaters_;
-  // The node position for each row, 1 HDV for each tree in the forest.
+  // The node position for each row, 1 HDV for each tree in the forest.  Note that the
+  // position is negated if the row is sampled out.
   std::vector<HostDeviceVector<bst_node_t>> node_position_;
   // Predictors
   std::unique_ptr<Predictor> cpu_predictor_;

From fd9fcc2b278510adf1ca6add76db5da21883d4e6 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 01:59:49 +0800
Subject: [PATCH 113/124] Cleanup.

---
 src/objective/regression_obj.cu                |  1 -
 tests/cpp/common/test_partition_builder.cc     |  4 +---
 tests/cpp/objective/test_regression_obj_gpu.cu | 10 ++--------
 tests/python-gpu/test_gpu_updaters.py          |  1 -
 4 files changed, 3 insertions(+), 13 deletions(-)

diff --git a/src/objective/regression_obj.cu b/src/objective/regression_obj.cu
index 7f81322f2929..3dc4a7b82316 100644
--- a/src/objective/regression_obj.cu
+++ b/src/objective/regression_obj.cu
@@ -5,7 +5,6 @@
  * \author Tianqi Chen, Kailong Chen
  */
 #include <dmlc/omp.h>
-#include <rabit/rabit.h>
 #include <xgboost/logging.h>
 #include <xgboost/objective.h>
 #include <xgboost/tree_model.h>
diff --git a/tests/cpp/common/test_partition_builder.cc b/tests/cpp/common/test_partition_builder.cc
index aeedf8e16b54..885b924e71c1 100644
--- a/tests/cpp/common/test_partition_builder.cc
+++ b/tests/cpp/common/test_partition_builder.cc
@@ -1,6 +1,3 @@
-/*!
- * Copyright 2021-2022 by XGBoost Contributors
- */
 #include <gtest/gtest.h>
 #include <vector>
 #include <string>
@@ -77,5 +74,6 @@ TEST(PartitionBuilder, BasicTest) {
     ASSERT_EQ(n_right, (kBlockSize - rows_for_left_node[nid]) * tasks[nid]);
   }
 }
+
 }  // namespace common
 }  // namespace xgboost
diff --git a/tests/cpp/objective/test_regression_obj_gpu.cu b/tests/cpp/objective/test_regression_obj_gpu.cu
index b6836207b5ef..38f29b8a8800 100644
--- a/tests/cpp/objective/test_regression_obj_gpu.cu
+++ b/tests/cpp/objective/test_regression_obj_gpu.cu
@@ -1,12 +1,6 @@
 /*!
- * Copyright 2018-2022 by XGBoost contributors
+ * Copyright 2018 XGBoost contributors
  */
-
+// Dummy file to keep the CUDA tests.
 
 #include "test_regression_obj.cc"
-
-namespace xgboost {
-namespace obj {
-
-}  // namespace obj
-}  // namespace xgboost
diff --git a/tests/python-gpu/test_gpu_updaters.py b/tests/python-gpu/test_gpu_updaters.py
index c6574bca4a22..a3427b566360 100644
--- a/tests/python-gpu/test_gpu_updaters.py
+++ b/tests/python-gpu/test_gpu_updaters.py
@@ -51,7 +51,6 @@ def test_gpu_hist(self, param, num_rounds, dataset):
         param["tree_method"] = "gpu_hist"
         param = dataset.set_params(param)
         result = train_result(param, dataset.get_dmat(), num_rounds)
-
         note(result)
         assert tm.non_increasing(result["train"][dataset.metric])
 

From 244c216e262f17281595abcd22021c320b2f5ad5 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 02:06:46 +0800
Subject: [PATCH 114/124] Skip updaters that don't support partitioner.

---
 include/xgboost/tree_updater.h   | 5 +++++
 src/gbm/gbtree.cc                | 3 +++
 src/tree/updater_approx.cc       | 2 ++
 src/tree/updater_gpu_hist.cu     | 2 ++
 src/tree/updater_quantile_hist.h | 2 ++
 5 files changed, 14 insertions(+)

diff --git a/include/xgboost/tree_updater.h b/include/xgboost/tree_updater.h
index 93cb14afa279..f0fabb26d9a0 100644
--- a/include/xgboost/tree_updater.h
+++ b/include/xgboost/tree_updater.h
@@ -49,6 +49,11 @@ class TreeUpdater : public Configurable {
    *  existing trees.
    */
   virtual bool CanModifyTree() const { return false; }
+  /*!
+   * \brief Wether the out_position in `Update` is valid. This determines whether adaptive
+   *        tree can be used.
+   */
+  virtual bool HasNodePosition() const { return false; }
   /*!
    * \brief perform update to the tree models
    * \param gpair the gradient pair statistics of the data
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index a96dccfc405b..a14c73d8615a 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -221,6 +221,9 @@ void GBTree::UpdateTreeLeaf(DMatrix const* p_fmat, HostDeviceVector<float> const
                             ObjFunction const* obj, size_t gidx,
                             std::vector<std::unique_ptr<RegTree>>* p_trees) {
   CHECK(!updaters_.empty());
+  if (!updaters_.back()->HasNodePosition()) {
+    return;
+  }
   if (!obj || !obj->Task().UpdateTreeLeaf()) {
     return;
   }
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index bf668b2a5e46..f5d1a8031214 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -355,6 +355,8 @@ class GlobalApproxUpdater : public TreeUpdater {
     }
     return true;
   }
+
+  bool HasNodePosition() const override { return true; }
 };
 
 DMLC_REGISTRY_FILE_TAG(grow_histmaker);
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 0bb6c601db80..0ee59db032ea 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -922,6 +922,8 @@ class GPUHistMaker : public TreeUpdater {
     return "grow_gpu_hist";
   }
 
+  bool HasNodePosition() const override { return true; }
+
  private:
   GPUHistMakerTrainParam hist_maker_param_;
   ObjInfo task_;
diff --git a/src/tree/updater_quantile_hist.h b/src/tree/updater_quantile_hist.h
index ec81f23ac5e1..8c6402eea249 100644
--- a/src/tree/updater_quantile_hist.h
+++ b/src/tree/updater_quantile_hist.h
@@ -276,6 +276,8 @@ class QuantileHistMaker: public TreeUpdater {
     return "grow_quantile_histmaker";
   }
 
+  bool HasNodePosition() const override { return true; }
+
  protected:
   CPUHistMakerTrainParam hist_maker_param_;
   // training parameter

From b940311b998c7ffff930510ebcbcbfa0fd0ba1d6 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 02:17:59 +0800
Subject: [PATCH 115/124] Forbid external memory.

---
 src/gbm/gbtree.cc         |  5 +++
 src/objective/adaptive.cc | 82 +++++++++++++++------------------------
 2 files changed, 36 insertions(+), 51 deletions(-)

diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index a14c73d8615a..bb7c341f8beb 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -249,6 +249,11 @@ void GBTree::DoBoost(DMatrix* p_fmat, HostDeviceVector<GradientPair>* in_gpair,
       {static_cast<size_t>(p_fmat->Info().num_row_), static_cast<size_t>(ngroup)},
       device};
   CHECK_NE(ngroup, 0);
+
+  if (!p_fmat->SingleColBlock() && obj->Task().UpdateTreeLeaf()) {
+    LOG(FATAL) << "Current objective doesn't support external memory.";
+  }
+
   if (ngroup == 1) {
     std::vector<std::unique_ptr<RegTree>> ret;
     BoostNewTrees(in_gpair, p_fmat, 0, &ret);
diff --git a/src/objective/adaptive.cc b/src/objective/adaptive.cc
index c43b6a1677d5..192ed785a0f6 100644
--- a/src/objective/adaptive.cc
+++ b/src/objective/adaptive.cc
@@ -84,61 +84,41 @@ void UpdateTreeLeafHost(Context const* ctx, std::vector<bst_node_t> const& posit
   std::vector<float> quantiles(n_leaf, 0);
   std::vector<int32_t> n_valids(n_leaf, 0);
 
-  {
-    std::vector<float> results(nidx.size());
-    auto const& h_node_idx = nidx;
-    auto const& h_node_ptr = nptr;
-    CHECK_LE(h_node_ptr.back(), info.num_row_);
-    // loop over each leaf
-    common::ParallelFor(results.size(), ctx->Threads(), [&](size_t k) {
-      auto nidx = h_node_idx[k];
-      CHECK(tree[nidx].IsLeaf());
-      CHECK_LT(k + 1, h_node_ptr.size());
-      size_t n = h_node_ptr[k + 1] - h_node_ptr[k];
-      auto h_row_set = common::Span<size_t const>{ridx}.subspan(h_node_ptr[k], n);
-      // multi-target not yet supported.
-      auto h_labels = info.labels.HostView().Slice(linalg::All(), 0);
-      auto const& h_predt = predt.ConstHostVector();
-      auto h_weights = linalg::MakeVec(&info.weights_);
-
-      auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
-        auto row_idx = h_row_set[i];
-        return h_labels(row_idx) - h_predt[row_idx];
-      });
-      auto w_it = common::MakeIndexTransformIter([&](size_t i) -> float {
-        auto row_idx = h_row_set[i];
-        return h_weights(row_idx);
-      });
-
-      float q{0};
-      if (info.weights_.Empty()) {
-        q = common::Quantile(alpha, iter, iter + h_row_set.size());
-      } else {
-        q = common::WeightedQuantile(alpha, iter, iter + h_row_set.size(), w_it);
-      }
-      if (std::isnan(q)) {
-        CHECK(h_row_set.empty());
-      }
-      results.at(k) = q;
+  auto const& h_node_idx = nidx;
+  auto const& h_node_ptr = nptr;
+  CHECK_LE(h_node_ptr.back(), info.num_row_);
+  // loop over each leaf
+  common::ParallelFor(quantiles.size(), ctx->Threads(), [&](size_t k) {
+    auto nidx = h_node_idx[k];
+    CHECK(tree[nidx].IsLeaf());
+    CHECK_LT(k + 1, h_node_ptr.size());
+    size_t n = h_node_ptr[k + 1] - h_node_ptr[k];
+    auto h_row_set = common::Span<size_t const>{ridx}.subspan(h_node_ptr[k], n);
+    // multi-target not yet supported.
+    auto h_labels = info.labels.HostView().Slice(linalg::All(), 0);
+    auto const& h_predt = predt.ConstHostVector();
+    auto h_weights = linalg::MakeVec(&info.weights_);
+
+    auto iter = common::MakeIndexTransformIter([&](size_t i) -> float {
+      auto row_idx = h_row_set[i];
+      return h_labels(row_idx) - h_predt[row_idx];
+    });
+    auto w_it = common::MakeIndexTransformIter([&](size_t i) -> float {
+      auto row_idx = h_row_set[i];
+      return h_weights(row_idx);
     });
 
-    // sum result from each external memory partition to quantiles
-    for (size_t i = 0; i < results.size(); ++i) {
-      if (!std::isnan(results[i])) {
-        quantiles[i] += results[i];
-        n_valids[i]++;
-      }
-    }
-  }
-
-  for (size_t i = 0; i < quantiles.size(); ++i) {
-    if (n_valids[i] > 0) {
-      quantiles[i] /= n_valids[i];
+    float q{0};
+    if (info.weights_.Empty()) {
+      q = common::Quantile(alpha, iter, iter + h_row_set.size());
     } else {
-      // mark that no page has valid sample in the i^th leaf
-      quantiles[i] = std::numeric_limits<float>::quiet_NaN();
+      q = common::WeightedQuantile(alpha, iter, iter + h_row_set.size(), w_it);
     }
-  }
+    if (std::isnan(q)) {
+      CHECK(h_row_set.empty());
+    }
+    quantiles.at(k) = q;
+  });
 
   UpdateLeafValues(&quantiles, nidx, p_tree);
 }

From fe1f36f3ca8f8559c3ba51aa94cd06ed93c97c47 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 02:57:51 +0800
Subject: [PATCH 116/124] Small cleanups.

---
 include/xgboost/objective.h  | 1 -
 include/xgboost/tree_model.h | 1 -
 src/common/common.h          | 3 +--
 3 files changed, 1 insertion(+), 4 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 52b87c5f2944..253ccb860d63 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -22,7 +22,6 @@
 
 namespace xgboost {
 
-struct RowIndexCache;
 class RegTree;
 
 /*! \brief interface of objective function */
diff --git a/include/xgboost/tree_model.h b/include/xgboost/tree_model.h
index 3658efd94ad5..b2d2ad3383de 100644
--- a/include/xgboost/tree_model.h
+++ b/include/xgboost/tree_model.h
@@ -23,7 +23,6 @@
 #include <algorithm>
 #include <tuple>
 #include <stack>
-#include "xgboost/host_device_vector.h"
 
 namespace xgboost {
 
diff --git a/src/common/common.h b/src/common/common.h
index d36d4a80ac0e..aa2d8197b4a1 100644
--- a/src/common/common.h
+++ b/src/common/common.h
@@ -1,5 +1,5 @@
 /*!
- * Copyright 2015-2018 by Contributors
+ * Copyright 2015-2022 by XGBoost Contributors
  * \file common.h
  * \brief Common utilities
  */
@@ -7,7 +7,6 @@
 #define XGBOOST_COMMON_COMMON_H_
 
 #include <xgboost/base.h>
-#include <xgboost/linalg.h>
 #include <xgboost/logging.h>
 #include <xgboost/span.h>
 

From a81359ffac10ddcb3a061e2c8fc0d4ab224cf877 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 04:03:40 +0800
Subject: [PATCH 117/124] Fix test.

---
 src/data/iterative_device_dmatrix.h   | 2 +-
 tests/python-gpu/test_gpu_updaters.py | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/data/iterative_device_dmatrix.h b/src/data/iterative_device_dmatrix.h
index ba2d4a92f9da..031b289f2760 100644
--- a/src/data/iterative_device_dmatrix.h
+++ b/src/data/iterative_device_dmatrix.h
@@ -68,7 +68,7 @@ class IterativeDeviceDMatrix : public DMatrix {
 
   BatchSet<EllpackPage> GetEllpackBatches(const BatchParam& param) override;
 
-  bool SingleColBlock() const override { return false; }
+  bool SingleColBlock() const override { return true; }
 
   MetaInfo &Info() override { return info_; }
   MetaInfo const &Info() const override { return info_; }
diff --git a/tests/python-gpu/test_gpu_updaters.py b/tests/python-gpu/test_gpu_updaters.py
index a3427b566360..e9d2bf06e229 100644
--- a/tests/python-gpu/test_gpu_updaters.py
+++ b/tests/python-gpu/test_gpu_updaters.py
@@ -90,6 +90,8 @@ def test_gpu_hist_device_dmatrix(self, param, num_rounds, dataset):
            tm.dataset_strategy)
     @settings(deadline=None, print_blob=True)
     def test_external_memory(self, param, num_rounds, dataset):
+        if dataset.name.endswith("-l1"):
+            return
         # We cannot handle empty dataset yet
         assume(len(dataset.y) > 0)
         param['tree_method'] = 'gpu_hist'

From 0d328f1031c0b680d3217d7ced8970d36e310ae5 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 17:43:21 +0800
Subject: [PATCH 118/124] tidy.

---
 src/objective/adaptive.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/objective/adaptive.cc b/src/objective/adaptive.cc
index 192ed785a0f6..f2675d918bdf 100644
--- a/src/objective/adaptive.cc
+++ b/src/objective/adaptive.cc
@@ -60,7 +60,6 @@ void EncodeTreeLeafHost(RegTree const& tree, std::vector<bst_node_t> const& posi
 
   if (n_leaf != leaf.size()) {
     FillMissingLeaf(leaf, &nidx, &nptr);
-    n_leaf = leaf.size();
   }
 }
 

From a0d883bca0bac0ebcc2ac34142c14300546fbbb6 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Sat, 23 Apr 2022 21:30:45 +0800
Subject: [PATCH 119/124] Use nan.

---
 src/objective/adaptive.h | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/objective/adaptive.h b/src/objective/adaptive.h
index 89d29a545cdc..85c041347cb9 100644
--- a/src/objective/adaptive.h
+++ b/src/objective/adaptive.h
@@ -4,6 +4,7 @@
 #pragma once
 
 #include <algorithm>
+#include <limits>
 #include <vector>
 
 #include "rabit/rabit.h"
@@ -40,7 +41,7 @@ inline void UpdateLeafValues(std::vector<float>* p_quantiles, std::vector<bst_no
   rabit::Allreduce<rabit::op::Max>(&n_leaf, 1);
   CHECK(quantiles.empty() || quantiles.size() == n_leaf);
   if (quantiles.empty()) {
-    quantiles.resize(n_leaf);
+    quantiles.resize(n_leaf, std::numeric_limits<float>::quiet_NaN());
   }
 
   // number of workers that have valid quantiles

From 5aa21bcaf39b277d14367ac30da8aa45e7847911 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Mon, 25 Apr 2022 19:53:15 +0800
Subject: [PATCH 120/124] Remove unused code.

---
 src/tree/updater_gpu_hist.cu | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 0ee59db032ea..bc8f14e8f010 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -682,7 +682,6 @@ struct GPUHistMakerDevice {
         int left_child_nidx = tree[candidate.nid].LeftChild();
         int right_child_nidx = tree[candidate.nid].RightChild();
         // Only create child entries if needed_
-        p_tree->GetNumLeaves();
         if (GPUExpandEntry::ChildIsValid(param, tree.GetDepth(left_child_nidx),
                                          num_leaves)) {
           monitor.Start("UpdatePosition");

From 1144fd3862e0ef581626e48f863c36330c4fe2f5 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Mon, 25 Apr 2022 20:14:30 +0800
Subject: [PATCH 121/124] Remove optimization.

---
 src/tree/gpu_hist/row_partitioner.cuh         | 20 ++--------
 src/tree/updater_gpu_hist.cu                  | 38 +++++++------------
 .../cpp/tree/gpu_hist/test_row_partitioner.cu |  9 +----
 3 files changed, 18 insertions(+), 49 deletions(-)

diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 52d3ca578e45..52a5fc5b6725 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -165,24 +165,10 @@ class RowPartitioner {
    *           argument and return the new position for this training instance.
    */
   template <typename FinalisePositionOpT, typename Sampledp>
-  void FinalisePosition(Context const* ctx, RegTree const* p_tree, size_t n_leaf, ObjInfo task,
-                        HostDeviceVector<bst_node_t>* p_out_row_indices, FinalisePositionOpT op,
-                        Sampledp sampledp) {
+  void FinalisePosition(Context const* ctx, HostDeviceVector<bst_node_t>* p_out_row_indices,
+                        FinalisePositionOpT op, Sampledp sampledp) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
-    if (!task.UpdateTreeLeaf()) {
-      dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
-        auto position = d_position[idx];
-        RowIndexT ridx = d_ridx[idx];
-        bst_node_t new_position = op(ridx, position);
-        if (new_position == kIgnoredTreePosition) {
-          return;
-        }
-        d_position[idx] = new_position;
-      });
-      return;
-    }
-
     p_out_row_indices->SetDevice(ctx->gpu_id);
     p_out_row_indices->Resize(position_.Size());
     auto sorted_position = p_out_row_indices->DevicePointer();
@@ -191,7 +177,7 @@ class RowPartitioner {
       RowIndexT ridx = d_ridx[idx];
       bst_node_t new_position = op(ridx, position);
       sorted_position[ridx] = sampledp(ridx) ? ~new_position : new_position;
-      if (new_position < 0) {
+      if (new_position == kIgnoredTreePosition) {
         return;
       }
       d_position[idx] = new_position;
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index bc8f14e8f010..56e34c4518e8 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -443,7 +443,7 @@ struct GPUHistMakerDevice {
     auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
     auto d_gpair = this->gpair;
     row_partitioner->FinalisePosition(
-        ctx_, p_tree, n_leaf, task, p_out_position,
+        ctx_, p_out_position,
         [=] __device__(size_t row_id, int position) {
           // What happens if user prune the tree?
           if (!d_matrix.IsInRange(row_id)) {
@@ -485,6 +485,7 @@ struct GPUHistMakerDevice {
   }
 
   void UpdatePredictionCache(linalg::VectorView<float> out_preds_d, RegTree const* p_tree) {
+    CHECK(p_tree);
     dh::safe_cuda(cudaSetDevice(ctx_->gpu_id));
     CHECK_EQ(out_preds_d.DeviceIdx(), ctx_->gpu_id);
     auto d_ridx = row_partitioner->GetRows();
@@ -499,29 +500,16 @@ struct GPUHistMakerDevice {
     auto d_node_sum_gradients = device_node_sum_gradients.data().get();
     auto tree_evaluator = evaluator_.GetEvaluator();
 
-    if (p_tree) {
-      auto const& h_nodes = p_tree->GetNodes();
-      dh::device_vector<RegTree::Node> nodes(h_nodes.size());
-      dh::safe_cuda(cudaMemcpyAsync(nodes.data().get(), h_nodes.data(),
-                                    h_nodes.size() * sizeof(RegTree::Node),
-                                    cudaMemcpyHostToDevice));
-      auto d_nodes = dh::ToSpan(nodes);
-      dh::LaunchN(d_ridx.size(), [=] XGBOOST_DEVICE(size_t idx) mutable {
-        bst_node_t nidx = d_position[idx];
-        auto weight = d_nodes[nidx].LeafValue();
-        out_preds_d(d_ridx[idx]) += weight;
-      });
-    } else {
-      // Avoid copying nodes by using evaluator to get leaf weight on-the-fly, this is
-      // useful when tree leaf is not updated after tree construction.
-      dh::LaunchN(d_ridx.size(), [=] XGBOOST_DEVICE(size_t local_idx) mutable {
-        bst_node_t nidx = d_position[local_idx];
-        float weight =
-            tree_evaluator.CalcWeight(nidx, param_d, GradStats{d_node_sum_gradients[nidx]});
-        static_assert(!std::is_const<decltype(out_preds_d)>::value, "");
-        out_preds_d(d_ridx[local_idx]) += weight * param_d.learning_rate;
-      });
-    }
+    auto const& h_nodes = p_tree->GetNodes();
+    dh::caching_device_vector<RegTree::Node> nodes(h_nodes.size());
+    dh::safe_cuda(cudaMemcpyAsync(nodes.data().get(), h_nodes.data(),
+                                  h_nodes.size() * sizeof(RegTree::Node), cudaMemcpyHostToDevice));
+    auto d_nodes = dh::ToSpan(nodes);
+    dh::LaunchN(d_ridx.size(), [=] XGBOOST_DEVICE(size_t idx) mutable {
+      bst_node_t nidx = d_position[idx];
+      auto weight = d_nodes[nidx].LeafValue();
+      out_preds_d(d_ridx[idx]) += weight;
+    });
     row_partitioner.reset();
   }
 
@@ -826,7 +814,7 @@ class GPUHistMakerSpecialised {
       return false;
     }
     monitor_.Start("UpdatePredictionCache");
-    maker->UpdatePredictionCache(p_out_preds, task_.zero_hess ? p_last_tree_ : nullptr);
+    maker->UpdatePredictionCache(p_out_preds, p_last_tree_);
     monitor_.Stop("UpdatePredictionCache");
     return true;
   }
diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index 126b51497f8f..9a2914f45293 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -109,9 +109,6 @@ TEST(RowPartitioner, Basic) { TestUpdatePosition(); }
 void TestFinalise() {
   const int kNumRows = 10;
 
-  ObjInfo task{ObjInfo::kRegression, false, false};
-  RegTree tree;
-  tree.ExpandNode(0, 0, 0.f, true, 0., 0., 0., /*loss_chg=*/0.f, 0.f, 0.f, 0.f);
   HostDeviceVector<bst_node_t> position;
   Context ctx;
   ctx.gpu_id = 0;
@@ -119,8 +116,7 @@ void TestFinalise() {
   {
     RowPartitioner rp(0, kNumRows);
     rp.FinalisePosition(
-        &ctx, &tree, tree.GetNumLeaves(), task, &position,
-        [=] __device__(RowPartitioner::RowIndexT ridx, int position) { return 7; },
+        &ctx, &position, [=] __device__(RowPartitioner::RowIndexT ridx, int position) { return 7; },
         [] XGBOOST_DEVICE(size_t idx) { return false; });
 
     auto position = rp.GetPositionHost();
@@ -143,11 +139,10 @@ void TestFinalise() {
   }
 
   auto d_hess = dh::ToSpan(hess);
-  task.zero_hess = true;
 
   RowPartitioner rp(0, kNumRows);
   rp.FinalisePosition(
-      &ctx, &tree, tree.GetNumLeaves(), task, &position,
+      &ctx, &position,
       [] __device__(RowPartitioner::RowIndexT ridx, bst_node_t position) {
         return ridx % 2 == 0 ? 1 : 2;
       },

From 8f96187aa10f7e61e723f2193f73b57cf475ad56 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Mon, 25 Apr 2022 21:33:35 +0800
Subject: [PATCH 122/124] Make it a check.

---
 src/common/partition_builder.h | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/common/partition_builder.h b/src/common/partition_builder.h
index 5e05a18d0ab4..648cbe61a3a3 100644
--- a/src/common/partition_builder.h
+++ b/src/common/partition_builder.h
@@ -291,9 +291,10 @@ class PartitionBuilder {
     auto p_begin = row_set.Data()->data();
     ParallelFor(row_set.Size(), ctx->Threads(), [&](size_t i) {
       auto const& node = row_set[i];
-      if (node.node_id < 0 || !tree[node.node_id].IsLeaf()) {
+      if (node.node_id < 0) {
         return;
       }
+      CHECK(tree[node.node_id].IsLeaf());
       if (node.begin) {  // guard for empty node.
         size_t ptr_offset = node.end - p_begin;
         CHECK_LE(ptr_offset, row_set.Data()->size()) << node.node_id;

From 1cf3f02399f496e88f1fc00d528c147283194c71 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Mon, 25 Apr 2022 22:12:15 +0800
Subject: [PATCH 123/124] Revert some cleanups.

---
 src/tree/gpu_hist/row_partitioner.cuh          | 18 ++++++++++++++++--
 src/tree/updater_gpu_hist.cu                   | 16 +++++++---------
 .../cpp/tree/gpu_hist/test_row_partitioner.cu  |  6 ++++--
 3 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 52a5fc5b6725..0afe5e5c678d 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -165,10 +165,24 @@ class RowPartitioner {
    *           argument and return the new position for this training instance.
    */
   template <typename FinalisePositionOpT, typename Sampledp>
-  void FinalisePosition(Context const* ctx, HostDeviceVector<bst_node_t>* p_out_row_indices,
-                        FinalisePositionOpT op, Sampledp sampledp) {
+  void FinalisePosition(Context const* ctx, ObjInfo task,
+                        HostDeviceVector<bst_node_t>* p_out_row_indices, FinalisePositionOpT op,
+                        Sampledp sampledp) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
+    if (!task.UpdateTreeLeaf()) {
+      dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
+        auto position = d_position[idx];
+        RowIndexT ridx = d_ridx[idx];
+        bst_node_t new_position = op(ridx, position);
+        if (new_position == kIgnoredTreePosition) {
+          return;
+        }
+        d_position[idx] = new_position;
+      });
+      return;
+    }
+
     p_out_row_indices->SetDevice(ctx->gpu_id);
     p_out_row_indices->Resize(position_.Size());
     auto sorted_position = p_out_row_indices->DevicePointer();
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 56e34c4518e8..569188fd5374 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -390,7 +390,7 @@ struct GPUHistMakerDevice {
   // After tree update is finished, update the position of all training
   // instances to their final leaf. This information is used later to update the
   // prediction cache
-  void FinalisePosition(RegTree const* p_tree, size_t n_leaf, DMatrix* p_fmat, ObjInfo task,
+  void FinalisePosition(RegTree const* p_tree, DMatrix* p_fmat, ObjInfo task,
                         HostDeviceVector<bst_node_t>* p_out_position) {
     dh::TemporaryArray<RegTree::Node> d_nodes(p_tree->GetNodes().size());
     dh::safe_cuda(cudaMemcpyAsync(d_nodes.data().get(), p_tree->GetNodes().data(),
@@ -419,31 +419,29 @@ struct GPUHistMakerDevice {
       LOG(FATAL) << "Current objective function can not be used with subsampled external memory.";
     }
     if (page->n_rows == p_fmat->Info().num_row_) {
-      FinalisePositionInPage(page, p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
+      FinalisePositionInPage(page, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
                              dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
-                             n_leaf, p_out_position);
+                             p_out_position);
     } else {
       for (auto const& batch : p_fmat->GetBatches<EllpackPage>(batch_param)) {
-        FinalisePositionInPage(batch.Impl(), p_tree, dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
+        FinalisePositionInPage(batch.Impl(), dh::ToSpan(d_nodes), dh::ToSpan(d_split_types),
                                dh::ToSpan(d_categories), dh::ToSpan(d_categories_segments), task,
-                               n_leaf, p_out_position);
+                               p_out_position);
       }
     }
   }
 
   void FinalisePositionInPage(EllpackPageImpl const *page,
-                              RegTree const* p_tree,
                               const common::Span<RegTree::Node> d_nodes,
                               common::Span<FeatureType const> d_feature_types,
                               common::Span<uint32_t const> categories,
                               common::Span<RegTree::Segment> categories_segments,
                               ObjInfo task,
-                              size_t n_leaf,
                               HostDeviceVector<bst_node_t>* p_out_position) {
     auto d_matrix = page->GetDeviceAccessor(ctx_->gpu_id);
     auto d_gpair = this->gpair;
     row_partitioner->FinalisePosition(
-        ctx_, p_out_position,
+        ctx_, task, p_out_position,
         [=] __device__(size_t row_id, int position) {
           // What happens if user prune the tree?
           if (!d_matrix.IsInRange(row_id)) {
@@ -699,7 +697,7 @@ struct GPUHistMakerDevice {
     }
 
     monitor.Start("FinalisePosition");
-    this->FinalisePosition(p_tree, num_leaves, p_fmat, task, p_out_position);
+    this->FinalisePosition(p_tree, p_fmat, task, p_out_position);
     monitor.Stop("FinalisePosition");
   }
 };
diff --git a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
index 9a2914f45293..c8aaf82dcb3e 100644
--- a/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
+++ b/tests/cpp/tree/gpu_hist/test_row_partitioner.cu
@@ -109,6 +109,7 @@ TEST(RowPartitioner, Basic) { TestUpdatePosition(); }
 void TestFinalise() {
   const int kNumRows = 10;
 
+  ObjInfo task{ObjInfo::kRegression, false, false};
   HostDeviceVector<bst_node_t> position;
   Context ctx;
   ctx.gpu_id = 0;
@@ -116,7 +117,8 @@ void TestFinalise() {
   {
     RowPartitioner rp(0, kNumRows);
     rp.FinalisePosition(
-        &ctx, &position, [=] __device__(RowPartitioner::RowIndexT ridx, int position) { return 7; },
+        &ctx, task, &position,
+        [=] __device__(RowPartitioner::RowIndexT ridx, int position) { return 7; },
         [] XGBOOST_DEVICE(size_t idx) { return false; });
 
     auto position = rp.GetPositionHost();
@@ -142,7 +144,7 @@ void TestFinalise() {
 
   RowPartitioner rp(0, kNumRows);
   rp.FinalisePosition(
-      &ctx, &position,
+      &ctx, task, &position,
       [] __device__(RowPartitioner::RowIndexT ridx, bst_node_t position) {
         return ridx % 2 == 0 ? 1 : 2;
       },

From a83c813e6125d7baaaec2b0beb720080ae403af8 Mon Sep 17 00:00:00 2001
From: jiamingy <jm.yuan@outlook.com>
Date: Mon, 25 Apr 2022 23:28:09 +0800
Subject: [PATCH 124/124] Naming after the refactor.

---
 include/xgboost/objective.h           |  3 ++-
 src/tree/gpu_hist/row_partitioner.cuh | 15 ++++++++-------
 src/tree/updater_approx.cc            |  4 ++--
 src/tree/updater_quantile_hist.cc     | 20 ++++++++++----------
 src/tree/updater_quantile_hist.h      |  6 +++---
 5 files changed, 25 insertions(+), 23 deletions(-)

diff --git a/include/xgboost/objective.h b/include/xgboost/objective.h
index 253ccb860d63..cb0fe7741dc9 100644
--- a/include/xgboost/objective.h
+++ b/include/xgboost/objective.h
@@ -98,8 +98,9 @@ class ObjFunction : public Configurable {
    *   computes only an average of quantile between workers. This breaks when some leaf
    *   have no sample assigned in a local worker.
    *
-   * \param row_index The index of rows for each output leaf.
+   * \param position The leaf index for each rows.
    * \param info MetaInfo providing labels and weights.
+   * \param prediction Model prediction after transformation.
    * \param p_tree Tree that needs to be updated.
    */
   virtual void UpdateTreeLeaf(HostDeviceVector<bst_node_t> const& position, MetaInfo const& info,
diff --git a/src/tree/gpu_hist/row_partitioner.cuh b/src/tree/gpu_hist/row_partitioner.cuh
index 0afe5e5c678d..9470b6447512 100644
--- a/src/tree/gpu_hist/row_partitioner.cuh
+++ b/src/tree/gpu_hist/row_partitioner.cuh
@@ -157,16 +157,17 @@ class RowPartitioner {
    * complete. Does not update any other meta information in this data structure, so
    * should only be used at the end of training.
    *
-   *   When the task requires update leaf, this function will copy the row partitions into
-   *   p_out_row_indices. Note that the node ptr might not start from 0 due to sampling.
+   *   When the task requires update leaf, this function will copy the node index into
+   *   p_out_position. The index is negated if it's being sampled in current iteration.
    *
-   * \param p_out_row_indices Row partitions for each leaf.
+   * \param p_out_position Node index for each row.
    * \param op Device lambda. Should provide the row index and current position as an
    *           argument and return the new position for this training instance.
+   * \param sampled A device lambda to inform the partitioner whether a row is sampled.
    */
   template <typename FinalisePositionOpT, typename Sampledp>
   void FinalisePosition(Context const* ctx, ObjInfo task,
-                        HostDeviceVector<bst_node_t>* p_out_row_indices, FinalisePositionOpT op,
+                        HostDeviceVector<bst_node_t>* p_out_position, FinalisePositionOpT op,
                         Sampledp sampledp) {
     auto d_position = position_.Current();
     const auto d_ridx = ridx_.Current();
@@ -183,9 +184,9 @@ class RowPartitioner {
       return;
     }
 
-    p_out_row_indices->SetDevice(ctx->gpu_id);
-    p_out_row_indices->Resize(position_.Size());
-    auto sorted_position = p_out_row_indices->DevicePointer();
+    p_out_position->SetDevice(ctx->gpu_id);
+    p_out_position->Resize(position_.Size());
+    auto sorted_position = p_out_position->DevicePointer();
     dh::LaunchN(position_.Size(), [=] __device__(size_t idx) {
       auto position = d_position[idx];
       RowIndexT ridx = d_ridx[idx];
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index f5d1a8031214..4222cddb1ee9 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -156,13 +156,13 @@ class GloablApproxBuilder {
   }
 
   void LeafPartition(RegTree const &tree, common::Span<float> hess,
-                     std::vector<bst_node_t> *p_out_row_indices) {
+                     std::vector<bst_node_t> *p_out_position) {
     monitor_->Start(__func__);
     if (!evaluator_.Task().UpdateTreeLeaf()) {
       return;
     }
     for (auto const &part : partitioner_) {
-      part.LeafPartition(ctx_, tree, hess, p_out_row_indices);
+      part.LeafPartition(ctx_, tree, hess, p_out_position);
     }
     monitor_->Stop(__func__);
   }
diff --git a/src/tree/updater_quantile_hist.cc b/src/tree/updater_quantile_hist.cc
index e870166217f0..011733b4582a 100644
--- a/src/tree/updater_quantile_hist.cc
+++ b/src/tree/updater_quantile_hist.cc
@@ -56,11 +56,11 @@ void QuantileHistMaker::Update(HostDeviceVector<GradientPair> *gpair, DMatrix *d
 
   size_t t_idx{0};
   for (auto p_tree : trees) {
-    auto &row_indices = out_position[t_idx];
+    auto &t_row_position = out_position[t_idx];
     if (hist_maker_param_.single_precision_histogram) {
-      this->float_builder_->UpdateTree(gpair, dmat, p_tree, &row_indices);
+      this->float_builder_->UpdateTree(gpair, dmat, p_tree, &t_row_position);
     } else {
-      this->double_builder_->UpdateTree(gpair, dmat, p_tree, &row_indices);
+      this->double_builder_->UpdateTree(gpair, dmat, p_tree, &t_row_position);
     }
     ++t_idx;
   }
@@ -176,13 +176,13 @@ void QuantileHistMaker::Builder<GradientSumT>::BuildHistogram(
 template <typename GradientSumT>
 void QuantileHistMaker::Builder<GradientSumT>::LeafPartition(
     RegTree const &tree, common::Span<GradientPair const> gpair,
-    std::vector<bst_node_t> *p_out_row_indices) {
+    std::vector<bst_node_t> *p_out_position) {
   monitor_->Start(__func__);
   if (!evaluator_->Task().UpdateTreeLeaf()) {
     return;
   }
   for (auto const &part : partitioner_) {
-    part.LeafPartition(ctx_, tree, gpair, p_out_row_indices);
+    part.LeafPartition(ctx_, tree, gpair, p_out_position);
   }
   monitor_->Stop(__func__);
 }
@@ -190,7 +190,7 @@ void QuantileHistMaker::Builder<GradientSumT>::LeafPartition(
 template <typename GradientSumT>
 void QuantileHistMaker::Builder<GradientSumT>::ExpandTree(
     DMatrix *p_fmat, RegTree *p_tree, const std::vector<GradientPair> &gpair_h,
-    HostDeviceVector<bst_node_t> *p_out_row_indices) {
+    HostDeviceVector<bst_node_t> *p_out_position) {
   monitor_->Start(__func__);
 
   Driver<CPUExpandEntry> driver(static_cast<TrainParam::TreeGrowPolicy>(param_.grow_policy));
@@ -247,15 +247,15 @@ void QuantileHistMaker::Builder<GradientSumT>::ExpandTree(
     expand_set = driver.Pop();
   }
 
-  auto &h_row_indices = p_out_row_indices->HostVector();
-  this->LeafPartition(tree, gpair_h, &h_row_indices);
+  auto &h_out_position = p_out_position->HostVector();
+  this->LeafPartition(tree, gpair_h, &h_out_position);
   monitor_->Stop(__func__);
 }
 
 template <typename GradientSumT>
 void QuantileHistMaker::Builder<GradientSumT>::UpdateTree(
     HostDeviceVector<GradientPair> *gpair, DMatrix *p_fmat, RegTree *p_tree,
-    HostDeviceVector<bst_node_t> *p_out_row_indices) {
+    HostDeviceVector<bst_node_t> *p_out_position) {
   monitor_->Start(__func__);
 
   std::vector<GradientPair> *gpair_ptr = &(gpair->HostVector());
@@ -268,7 +268,7 @@ void QuantileHistMaker::Builder<GradientSumT>::UpdateTree(
 
   this->InitData(p_fmat, *p_tree, gpair_ptr);
 
-  ExpandTree(p_fmat, p_tree, *gpair_ptr, p_out_row_indices);
+  ExpandTree(p_fmat, p_tree, *gpair_ptr, p_out_position);
   monitor_->Stop(__func__);
 }
 
diff --git a/src/tree/updater_quantile_hist.h b/src/tree/updater_quantile_hist.h
index 8c6402eea249..6d5919abb75f 100644
--- a/src/tree/updater_quantile_hist.h
+++ b/src/tree/updater_quantile_hist.h
@@ -302,7 +302,7 @@ class QuantileHistMaker: public TreeUpdater {
     }
     // update one tree, growing
     void UpdateTree(HostDeviceVector<GradientPair>* gpair, DMatrix* p_fmat, RegTree* p_tree,
-                    HostDeviceVector<bst_node_t>* p_out_row_indices);
+                    HostDeviceVector<bst_node_t>* p_out_position);
 
     bool UpdatePredictionCache(DMatrix const* data, linalg::VectorView<float> out_preds) const;
 
@@ -322,10 +322,10 @@ class QuantileHistMaker: public TreeUpdater {
                         std::vector<GradientPair> const& gpair);
 
     void LeafPartition(RegTree const& tree, common::Span<GradientPair const> gpair,
-                       std::vector<bst_node_t>* p_out_row_indices);
+                       std::vector<bst_node_t>* p_out_position);
 
     void ExpandTree(DMatrix* p_fmat, RegTree* p_tree, const std::vector<GradientPair>& gpair_h,
-                    HostDeviceVector<bst_node_t>* p_out_row_indices);
+                    HostDeviceVector<bst_node_t>* p_out_position);
 
    private:
     const size_t n_trees_;