Support leaf prediction with QDM on CPU. (#11620)

Author: trivialfis

python-package/xgboost/testing/dask.py

@@ -322,16 +322,18 @@ def pack(**kwargs: Any) -> dd.DataFrame:
     return X, y
 
 
+# pylint: disable=too-many-locals
 def run_recode(client: Client, device: Device) -> None:
     """Run re-coding test with the Dask interface."""
     enc, reenc, y, _, _ = make_recoded(device, n_features=96)
+    workers = get_client_workers(client)
     denc, dreenc, dy = (
-        dd.from_pandas(enc, npartitions=8),
-        dd.from_pandas(reenc, npartitions=8),
-        da.from_array(y, chunks=(y.shape[0] // 8,)),
+        dd.from_pandas(enc, npartitions=8).persist(workers=workers),
+        dd.from_pandas(reenc, npartitions=8).persist(workers=workers),
+        da.from_array(y, chunks=(y.shape[0] // 8,)).persist(workers=workers),
     )
+
     wait([denc, dreenc, dy])
-    client.rebalance([denc, dreenc, dy])
 
     if device == "cuda":
         denc = denc.to_backend("cudf")

python-package/xgboost/testing/predict.py

@@ -0,0 +1,65 @@
+"""Tests for inference."""
+
+from typing import Type
+
+import numpy as np
+
+from ..core import DMatrix
+from ..training import train
+from .shared import validate_leaf_output
+from .utils import Device
+
+
+# pylint: disable=invalid-name,too-many-locals
+def run_predict_leaf(device: Device, DMatrixT: Type[DMatrix]) -> np.ndarray:
+    """Run tests for leaf index prediction."""
+    rows = 100
+    cols = 4
+    classes = 5
+    num_parallel_tree = 4
+    num_boost_round = 10
+    rng = np.random.RandomState(1994)
+    X = rng.randn(rows, cols)
+    y = rng.randint(low=0, high=classes, size=rows)
+
+    m = DMatrixT(X, y)
+    booster = train(
+        {
+            "num_parallel_tree": num_parallel_tree,
+            "num_class": classes,
+            "tree_method": "hist",
+        },
+        m,
+        num_boost_round=num_boost_round,
+    )
+
+    booster.set_param({"device": device})
+    empty = DMatrixT(np.ones(shape=(0, cols)))
+    empty_leaf = booster.predict(empty, pred_leaf=True)
+    assert empty_leaf.shape[0] == 0
+
+    leaf = booster.predict(m, pred_leaf=True, strict_shape=True)
+    assert leaf.shape[0] == rows
+    assert leaf.shape[1] == num_boost_round
+    assert leaf.shape[2] == classes
+    assert leaf.shape[3] == num_parallel_tree
+
+    validate_leaf_output(leaf, num_parallel_tree)
+
+    n_iters = np.int32(2)
+    sliced = booster.predict(
+        m,
+        pred_leaf=True,
+        iteration_range=(0, n_iters),
+        strict_shape=True,
+    )
+    first = sliced[0, ...]
+
+    assert np.prod(first.shape) == classes * num_parallel_tree * n_iters
+
+    # When there's only 1 tree, the output is a 1 dim vector
+    booster = train({"tree_method": "hist"}, num_boost_round=1, dtrain=m)
+    booster.set_param({"device": device})
+    assert booster.predict(m, pred_leaf=True).shape == (rows,)
+
+    return leaf

src/common/threading_utils.h

@@ -251,6 +251,17 @@ void ParallelFor(Index size, std::int32_t n_threads, Func&& fn) {
   ParallelFor(size, n_threads, Sched::Static(), std::forward<Func>(fn));
 }
 
+/**
+ * @brief 1-d block-based parallel for loop.
+ *
+ * @tparam kBlockOfRowsSize The size of the block.
+ * @tparam Index The type of the index.
+ * @tparam Func The type of the function.
+ *
+ * @param size The size of the range.
+ * @param n_threads The number of threads.
+ * @param fn The function to execute. The function should take a Range1d as an argument.
+ */
 template <std::size_t kBlockOfRowsSize, typename Index, typename Func>
 void ParallelFor1d(Index size, std::int32_t n_threads, Func&& fn) {
   static_assert(std::is_void_v<std::invoke_result_t<Func, common::Range1d>>);

src/data/ellpack_page.cu

@@ -462,6 +462,10 @@ EllpackPageImpl::EllpackPageImpl(Context const* ctx, GHistIndexMatrix const& pag
       info{CalcNumSymbols(
           ctx,
           [&] {
+            if (page.Size() == 0) {
+              return static_cast<typename decltype(page.row_ptr)::value_type>(0);
+            }
+            CHECK_GE(page.row_ptr.size(), 2);
             auto it = common::MakeIndexTransformIter(
                 [&](bst_idx_t i) { return page.row_ptr[i + 1] - page.row_ptr[i]; });
             return *std::max_element(it, it + page.Size());

src/data/gradient_index.h

@@ -20,7 +20,7 @@
 #include "../common/ref_resource_view.h"  // for RefResourceView
 #include "../common/threading_utils.h"
 #include "../common/transform_iterator.h"  // for MakeIndexTransformIter
-#include "adapter.h"
+#include "entry.h"                         // for IsValidFunctor
 #include "xgboost/base.h"
 #include "xgboost/data.h"
 

src/data/proxy_dmatrix.h

@@ -158,6 +158,7 @@ struct ExternalDataInfo {
 
     CHECK_GE(this->n_features, 1) << "Data must has at least 1 column.";
     CHECK_EQ(this->base_rowids.size(), this->n_batches + 1);
+    CHECK_LE(this->row_stride, this->n_features);
   }
 
   void SetInfo(Context const* ctx, bool sync, MetaInfo* p_info) {