Use predictor for leaf id cache updates

include/xgboost/predictor.h

@@ -12,6 +12,8 @@
 #include <xgboost/context.h>
 #include <xgboost/data.h>
 #include <xgboost/host_device_vector.h>
+#include <xgboost/linalg.h>
+#include <xgboost/span.h>
 
 #include <functional>  // for function
 #include <memory>      // for shared_ptr
@@ -24,6 +26,7 @@ struct GBTreeModel;
 }  // namespace xgboost::gbm
 
 namespace xgboost {
+class RegTree;
 /**
  * \brief Contains pointer to input matrix and associated cached predictions.
  */
@@ -141,6 +144,18 @@ class Predictor {
   virtual void PredictLeaf(DMatrix* dmat, HostDeviceVector<float>* out_preds,
                            gbm::GBTreeModel const& model, bst_tree_t tree_end = 0) const = 0;
 
+  /**
+   * \brief Add prediction contributions from known leaf ids. The leaf ids are encoded with
+   *        tree::SamplePosition, where invalid rows are still decoded for prediction.
+   *
+   * \param leaf_ids Leaf ids for each tree, one vector per tree.
+   * \param trees Trees corresponding to the leaf id vectors.
+   * \param out_preds Prediction output to be incremented.
+   */
+  virtual void PredictFromLeafIds(common::Span<HostDeviceVector<bst_node_t> const> leaf_ids,
+                                  common::Span<RegTree const*> trees,
+                                  linalg::MatrixView<float> out_preds) const = 0;
+
   /**
    * \brief feature contributions to individual predictions; the output will be
    * a vector of length (nfeats + 1) * num_output_group * nsample, arranged in

include/xgboost/tree_updater.h

@@ -72,24 +72,6 @@ class TreeUpdater : public Configurable {
                       common::Span<HostDeviceVector<bst_node_t>> out_position,
                       std::vector<RegTree*> const& out_trees) = 0;
 
-  /**
-   * @brief Determines whether updater has enough knowledge about a given dataset to
-   *        quickly update prediction cache for the training data and performs the update
-   *        if possible.
-   *
-   * @param p_fmat data matrix
-   * @param out_preds prediction cache to be updated
-   *
-   * @return boolean indicating whether updater has capability to update the prediction
-   *         cache. If true, the prediction cache will have been updated by the time this
-   *         function returns.
-   */
-  virtual bool UpdatePredictionCache(DMatrix const* /*data*/,
-                                     common::Span<HostDeviceVector<bst_node_t>> /*out_position*/,
-                                     linalg::MatrixView<float> /*out_preds*/) {
-    return false;
-  }
-
   [[nodiscard]] virtual char const* Name() const = 0;
 
   /**

plugin/sycl/predictor/predictor.cc

@@ -10,8 +10,8 @@
 
 #include <cstddef>
 #include <limits>
-#include <mutex>
 #include <sycl/sycl.hpp>
+#include <vector>
 
 #include "../../../src/common/timer.h"
 #include "../data.h"
@@ -25,6 +25,7 @@
 #pragma GCC diagnostic pop
 #include "../../src/common/math.h"
 #include "../../src/gbm/gbtree_model.h"
+#include "../../src/tree/sample_position.h"
 #include "../device_manager.h"
 #include "../device_properties.h"
 #include "node.h"
@@ -217,6 +218,55 @@ class Predictor : public xgboost::Predictor {
     cpu_predictor->PredictLeaf(p_fmat, out_preds, model, ntree_limit);
   }
 
+  void PredictFromLeafIds(common::Span<HostDeviceVector<bst_node_t> const> leaf_ids,
+                          common::Span<RegTree const*> trees,
+                          linalg::MatrixView<float> out_preds) const override {
+    CHECK_EQ(leaf_ids.size(), trees.size());
+    if (out_preds.Device().IsCPU()) {
+      cpu_predictor->PredictFromLeafIds(leaf_ids, trees, out_preds);
+      return;
+    }
+    CHECK(out_preds.Device().IsSycl());
+    qu_ = device_manager.GetQueue(out_preds.Device());
+
+    for (std::size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
+      auto const* p_tree = trees[tree_idx];
+      CHECK(p_tree);
+      CHECK(!p_tree->IsMultiTarget())
+          << "Multi-target leaf-id prediction is not implemented for SYCL.";
+      CHECK_EQ(out_preds.Shape(1), 1);
+
+      auto d_leaf_ids = leaf_ids[tree_idx].ConstDeviceSpan();
+      CHECK_EQ(d_leaf_ids.size(), out_preds.Shape(0));
+
+      auto const h_nodes = p_tree->GetNodes(DeviceOrd::CPU());
+      std::vector<float> h_leaf_values(h_nodes.size());
+      for (std::size_t nidx = 0; nidx < h_nodes.size(); ++nidx) {
+        if (h_nodes[nidx].IsLeaf()) {
+          h_leaf_values[nidx] = h_nodes[nidx].LeafValue();
+        }
+      }
+      USMVector<float, MemoryType::on_device> leaf_values;
+      leaf_values.Resize(qu_, h_leaf_values.size());
+      qu_->memcpy(leaf_values.Data(), h_leaf_values.data(),
+                  h_leaf_values.size() * sizeof(h_leaf_values.front()))
+          .wait_and_throw();
+
+      auto const* d_leaf_values = leaf_values.DataConst();
+      auto* out = out_preds.Values().data();
+      auto out_stride = out_preds.Stride(0);
+      qu_->submit([&](::sycl::handler& cgh) {
+           cgh.parallel_for<>(::sycl::range<1>(d_leaf_ids.size()), [=](::sycl::id<1> pid) {
+             auto row_idx = pid[0];
+             auto position = d_leaf_ids[row_idx];
+             auto nidx = position >= 0 ? position : ~position;
+             out[row_idx * out_stride] += d_leaf_values[nidx];
+           });
+         })
+          .wait_and_throw();
+    }
+  }
+
   void PredictContribution(DMatrix* p_fmat, HostDeviceVector<float>* out_contribs,
                            const gbm::GBTreeModel& model, bst_tree_t ntree_limit, bool approximate,
                            int condition, unsigned condition_feature) const override {

plugin/sycl/tree/hist_updater.cc

@@ -11,6 +11,7 @@
 #include "../../src/collective/allreduce.h"
 #include "../../src/tree/common_row_partitioner.h"
 #include "../../src/tree/driver.h"
+#include "../../src/tree/sample_position.h"
 #include "../common/hist_util.h"
 #include "xgboost/linalg.h"
 
@@ -309,10 +310,11 @@ void HistUpdater<GradientSumT>::ExpandWithLossGuide(const common::GHistIndexMatr
 }
 
 template <typename GradientSumT>
-void HistUpdater<GradientSumT>::Update(
-    xgboost::tree::TrainParam const* param, const common::GHistIndexMatrix& gmat,
-    const HostDeviceVector<GradientPair>& gpair, DMatrix* p_fmat,
-    xgboost::common::Span<HostDeviceVector<bst_node_t>> out_position, RegTree* p_tree) {
+void HistUpdater<GradientSumT>::Update(xgboost::tree::TrainParam const* param,
+                                       const common::GHistIndexMatrix& gmat,
+                                       const HostDeviceVector<GradientPair>& gpair, DMatrix* p_fmat,
+                                       HostDeviceVector<bst_node_t>* p_out_position,
+                                       RegTree* p_tree) {
   builder_monitor_.Start("Update");
 
   tree_evaluator_.Reset(qu_, param_, p_fmat->Info().num_col_);
@@ -330,54 +332,42 @@ void HistUpdater<GradientSumT>::Update(
     p_tree->Stat(nid).base_weight = snode_host_[nid].weight;
     p_tree->Stat(nid).sum_hess = static_cast<float>(snode_host_[nid].stats.GetHess());
   }
+  this->FinalizePosition(p_fmat->Info().num_row_, *p_tree, p_out_position);
 
   builder_monitor_.Stop("Update");
 }
 
 template <typename GradientSumT>
-bool HistUpdater<GradientSumT>::UpdatePredictionCache(
-    const DMatrix* data, ::xgboost::linalg::MatrixView<float> out_preds) {
-  CHECK(out_preds.Device().IsSycl());
-  // p_last_fmat_ is a valid pointer as long as UpdatePredictionCache() is called in
-  // conjunction with Update().
-  if (!p_last_fmat_ || !p_last_tree_ || data != p_last_fmat_) {
-    return false;
+void HistUpdater<GradientSumT>::FinalizePosition(std::size_t n_samples, RegTree const& tree,
+                                                 HostDeviceVector<bst_node_t>* p_out_position) {
+  CHECK(p_out_position);
+  if (param_.subsample < 1.0f || row_set_collection_.Data().Size() != n_samples) {
+    p_out_position->Resize(0);
+    return;
   }
-  builder_monitor_.Start("UpdatePredictionCache");
-  CHECK_GT(out_preds.Size(), 0U);
-
-  size_t n_nodes = row_set_collection_.Size();
-  std::vector<::sycl::event> events(n_nodes);
-  auto tree = p_last_tree_->HostScView();
-  for (size_t node = 0; node < n_nodes; node++) {
-    const common::RowSetCollection::Elem& rowset = row_set_collection_[node];
-    if (rowset.begin != nullptr && rowset.end != nullptr && rowset.Size() != 0) {
-      int nid = rowset.node_id;
-      // if a node is marked as deleted by the pruner, traverse upward to locate
-      // a non-deleted leaf.
-      if (tree.IsDeleted(nid)) {
-        while (tree.IsDeleted(nid)) {
-          nid = tree.Parent(nid);
-        }
-        CHECK(tree.IsLeaf(nid));
-      }
-      bst_float leaf_value = tree.LeafValue(nid);
-      const size_t* rid = rowset.begin;
-      const size_t num_rows = rowset.Size();
-
-      events[node] = qu_->submit([&](::sycl::handler& cgh) {
-        cgh.parallel_for<>(::sycl::range<1>(num_rows), [=](::sycl::item<1> pid) {
-          size_t row_id = rid[pid.get_id(0)];
-          float& val = const_cast<float&>(out_preds(row_id));
-          val += leaf_value;
-        });
-      });
-    }
+  p_out_position->SetDevice(ctx_->Device());
+  p_out_position->Resize(n_samples);
+  if (n_samples == 0) {
+    return;
   }
-  qu_->wait();
+  auto d_position = p_out_position->DeviceSpan();
 
-  builder_monitor_.Stop("UpdatePredictionCache");
-  return true;
+  auto const t = tree.HostScView();
+  for (std::size_t nid = 0; nid < row_set_collection_.Size(); ++nid) {
+    auto const& row_set = row_set_collection_[nid];
+    if (row_set.node_id < 0 || !t.IsLeaf(row_set.node_id) || row_set.Size() == 0) {
+      continue;
+    }
+    auto const encoded = xgboost::tree::SamplePosition::Encode(row_set.node_id, true);
+    auto const* rows = row_set.begin;
+    auto* position = d_position.data();
+    auto n_rows = row_set.Size();
+    qu_->submit([&](::sycl::handler& cgh) {
+         cgh.parallel_for<>(::sycl::range<1>(n_rows),
+                            [=](::sycl::id<1> pid) { position[rows[pid[0]]] = encoded; });
+       })
+        .wait_and_throw();
+  }
 }
 
 template <typename GradientSumT>

plugin/sycl/tree/hist_updater.h

@@ -77,9 +77,7 @@ class HistUpdater {
   // update one tree, growing
   void Update(xgboost::tree::TrainParam const* param, const common::GHistIndexMatrix& gmat,
               const HostDeviceVector<GradientPair>& gpair, DMatrix* p_fmat,
-              xgboost::common::Span<HostDeviceVector<bst_node_t>> out_position, RegTree* p_tree);
-
-  bool UpdatePredictionCache(const DMatrix* data, ::xgboost::linalg::MatrixView<float> p_out_preds);
+              HostDeviceVector<bst_node_t>* p_out_position, RegTree* p_tree);
 
   void SetHistSynchronizer(HistSynchronizer<GradientSumT>* sync);
   void SetHistRowsAdder(HistRowsAdder<GradientSumT>* adder);
@@ -164,6 +162,9 @@ class HistUpdater {
   void ExpandWithLossGuide(const common::GHistIndexMatrix& gmat, RegTree* p_tree,
                            const HostDeviceVector<GradientPair>& gpair);
 
+  void FinalizePosition(std::size_t n_samples, RegTree const& tree,
+                        HostDeviceVector<bst_node_t>* p_out_position);
+
   void ReduceHists(const std::vector<int>& sync_ids, size_t nbins);
 
   inline static bool LossGuide(ExpandEntry lhs, ExpandEntry rhs) {

plugin/sycl/tree/updater_quantile_hist.cc

@@ -2,8 +2,8 @@
  * Copyright 2017-2024 by Contributors
  * \file updater_quantile_hist.cc
  */
-#include <vector>
 #include <memory>
+#include <vector>
 
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wtautological-constant-compare"
@@ -25,7 +25,7 @@ DMLC_REGISTRY_FILE_TAG(updater_quantile_hist_sycl);
 
 DMLC_REGISTER_PARAMETER(HistMakerTrainParam);
 
-void QuantileHistMaker::Configure(const Args& args) {
+void QuantileHistMaker::Configure(const Args &args) {
   const DeviceOrd device_spec = ctx_->Device();
   qu_ = device_manager.GetQueue(device_spec);
 
@@ -43,14 +43,9 @@ void QuantileHistMaker::Configure(const Args& args) {
   }
 }
 
-template<typename GradientSumT>
-void QuantileHistMaker::SetPimpl(std::unique_ptr<HistUpdater<GradientSumT>>* pimpl,
-                                 DMatrix *dmat) {
-  pimpl->reset(new HistUpdater<GradientSumT>(
-                ctx_,
-                qu_,
-                param_,
-                int_constraint_, dmat));
+template <typename GradientSumT>
+void QuantileHistMaker::SetPimpl(std::unique_ptr<HistUpdater<GradientSumT>> *pimpl, DMatrix *dmat) {
+  pimpl->reset(new HistUpdater<GradientSumT>(ctx_, qu_, param_, int_constraint_, dmat));
   if (collective::IsDistributed()) {
     (*pimpl)->SetHistSynchronizer(new DistributedHistSynchronizer<GradientSumT>());
     (*pimpl)->SetHistRowsAdder(new DistributedHistRowsAdder<GradientSumT>());
@@ -66,8 +61,9 @@ void QuantileHistMaker::CallUpdate(const std::unique_ptr<HistUpdater<GradientSum
                                    ::xgboost::linalg::Matrix<GradientPair> *gpair, DMatrix *dmat,
                                    xgboost::common::Span<HostDeviceVector<bst_node_t>> out_position,
                                    const std::vector<RegTree *> &trees) {
-  for (auto tree : trees) {
-    pimpl->Update(param, gmat_, *(gpair->Data()), dmat, out_position, tree);
+  CHECK_EQ(out_position.size(), trees.size());
+  for (std::size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
+    pimpl->Update(param, gmat_, *(gpair->Data()), dmat, &out_position[tree_idx], trees[tree_idx]);
   }
 }
 
@@ -105,30 +101,9 @@ void QuantileHistMaker::Update(xgboost::tree::TrainParam const *param, GradientC
   p_last_dmat_ = dmat;
 }
 
-bool QuantileHistMaker::UpdatePredictionCache(const DMatrix *data,
-                                              xgboost::common::Span<HostDeviceVector<bst_node_t>>,
-                                              ::xgboost::linalg::MatrixView<float> out_preds) {
-  if (param_.subsample < 1.0f) return false;
-
-  if (hist_precision_ == HistPrecision::fp32) {
-    if (pimpl_fp32) {
-      return pimpl_fp32->UpdatePredictionCache(data, out_preds);
-    } else {
-      return false;
-    }
-  } else {
-    if (pimpl_fp64) {
-      return pimpl_fp64->UpdatePredictionCache(data, out_preds);
-    } else {
-      return false;
-    }
-  }
-}
-
 XGBOOST_REGISTER_TREE_UPDATER(QuantileHistMaker, "grow_quantile_histmaker_sycl")
-.describe("Grow tree using quantized histogram with SYCL.")
-.set_body(
-    [](Context const* ctx, ObjInfo const * task) {
+    .describe("Grow tree using quantized histogram with SYCL.")
+    .set_body([](Context const *ctx, ObjInfo const *task) {
       return new QuantileHistMaker(ctx, task);
     });
 }  // namespace tree