Predict on Ellpack. (#5327)

* Unify GPU prediction node.
* Add `PageExists`.
* Dispatch prediction on input data for GPU Predictor.

tests/cpp/predictor/test_gpu_predictor.cu

@@ -11,11 +11,12 @@
 #include "gtest/gtest.h"
 #include "../helpers.h"
 #include "../../../src/gbm/gbtree_model.h"
+#include "test_predictor.h"
 
 namespace xgboost {
 namespace predictor {
 
-TEST(GpuPredictor, Basic) {
+TEST(GPUPredictor, Basic) {
   auto cpu_lparam = CreateEmptyGenericParam(-1);
   auto gpu_lparam = CreateEmptyGenericParam(0);
 
@@ -56,7 +57,20 @@ TEST(GpuPredictor, Basic) {
   }
 }
 
-TEST(gpu_predictor, ExternalMemoryTest) {
+TEST(GPUPredictor, EllpackBasic) {
+  for (size_t bins = 2; bins < 258; bins += 16) {
+    size_t rows = bins * 16;
+    TestPredictionFromGradientIndex<EllpackPage>("gpu_predictor", rows, bins);
+    TestPredictionFromGradientIndex<EllpackPage>("gpu_predictor", bins, bins);
+  }
+}
+
+TEST(GPUPredictor, EllpackTraining) {
+  size_t constexpr kRows { 128 };
+  TestTrainingPrediction(kRows, "gpu_hist");
+}
+
+TEST(GPUPredictor, ExternalMemoryTest) {
   auto lparam = CreateEmptyGenericParam(0);
   std::unique_ptr<Predictor> gpu_predictor =
       std::unique_ptr<Predictor>(Predictor::Create("gpu_predictor", &lparam));

tests/cpp/predictor/test_predictor.cc

@@ -2,13 +2,16 @@
  * Copyright 2020 by Contributors
  */
 
-#include <cstddef>
 #include <gtest/gtest.h>
 #include <xgboost/predictor.h>
 #include <xgboost/data.h>
+#include <xgboost/host_device_vector.h>
+#include <xgboost/generic_parameters.h>
+
+#include "test_predictor.h"
 
 #include "../helpers.h"
-#include "xgboost/generic_parameters.h"
+#include "../../../src/common/io.h"
 
 namespace xgboost {
 TEST(Predictor, PredictionCache) {
@@ -30,4 +33,52 @@ TEST(Predictor, PredictionCache) {
   add_cache();
   EXPECT_ANY_THROW(container.Entry(m));
 }
+
+// Only run this test when CUDA is enabled.
+void TestTrainingPrediction(size_t rows, std::string tree_method) {
+  size_t constexpr kCols = 16;
+  size_t constexpr kClasses = 3;
+  size_t constexpr kIters = 3;
+
+  std::unique_ptr<Learner> learner;
+  auto train = [&](std::string predictor, HostDeviceVector<float>* out) {
+    auto pp_m = CreateDMatrix(rows, kCols, 0);
+    auto p_m = *pp_m;
+
+    auto &h_label = p_m->Info().labels_.HostVector();
+    h_label.resize(rows);
+
+    for (size_t i = 0; i < rows; ++i) {
+      h_label[i] = i % kClasses;
+    }
+
+    learner.reset(Learner::Create({}));
+    learner->SetParam("tree_method", tree_method);
+    learner->SetParam("objective", "multi:softprob");
+    learner->SetParam("predictor", predictor);
+    learner->SetParam("num_feature", std::to_string(kCols));
+    learner->SetParam("num_class", std::to_string(kClasses));
+    learner->Configure();
+
+    for (size_t i = 0; i < kIters; ++i) {
+      learner->UpdateOneIter(i, p_m);
+    }
+    learner->Predict(p_m, false, out);
+    delete pp_m;
+  };
+  // Alternate the predictor, CPU predictor can not use ellpack while GPU predictor can
+  // not use CPU histogram index.  So it's guaranteed one of the following is not
+  // predicting from histogram index.  Note: As of writing only GPU supports predicting
+  // from gradient index, the test is written for future portability.
+  HostDeviceVector<float> predictions_0;
+  train("cpu_predictor", &predictions_0);
+
+  HostDeviceVector<float> predictions_1;
+  train("gpu_predictor", &predictions_1);
+
+  for (size_t i = 0; i < rows; ++i) {
+    EXPECT_NEAR(predictions_1.ConstHostVector()[i],
+                predictions_0.ConstHostVector()[i], kRtEps);
+  }
+}
 }  // namespace xgboost

tests/cpp/predictor/test_predictor.h

@@ -0,0 +1,70 @@
+#ifndef XGBOOST_TEST_PREDICTOR_H_
+#define XGBOOST_TEST_PREDICTOR_H_
+
+#include <xgboost/predictor.h>
+#include <string>
+#include <cstddef>
+#include "../helpers.h"
+
+namespace xgboost {
+template <typename Page>
+void TestPredictionFromGradientIndex(std::string name, size_t rows, int32_t bins) {
+  constexpr size_t kCols { 8 }, kClasses { 3 };
+
+  LearnerModelParam param;
+  param.num_feature = kCols;
+  param.num_output_group = kClasses;
+  param.base_score = 0.5;
+
+  auto lparam = CreateEmptyGenericParam(0);
+
+  std::unique_ptr<Predictor> predictor =
+      std::unique_ptr<Predictor>(Predictor::Create(name, &lparam));
+  predictor->Configure({});
+
+  gbm::GBTreeModel model = CreateTestModel(&param, kClasses);
+
+  {
+    auto pp_ellpack = CreateDMatrix(rows, kCols, 0);
+    auto p_ellpack = *pp_ellpack;
+    // Use same number of bins as rows.
+    for (auto const &page DMLC_ATTRIBUTE_UNUSED :
+         p_ellpack->GetBatches<Page>({0, static_cast<int32_t>(bins), 0})) {
+    }
+
+    auto pp_precise = CreateDMatrix(rows, kCols, 0);
+    auto p_precise = *pp_precise;
+
+    PredictionCacheEntry approx_out_predictions;
+    predictor->PredictBatch(p_ellpack.get(), &approx_out_predictions, model, 0);
+
+    PredictionCacheEntry precise_out_predictions;
+    predictor->PredictBatch(p_precise.get(), &precise_out_predictions, model, 0);
+
+    for (size_t i = 0; i < rows; ++i) {
+      CHECK_EQ(approx_out_predictions.predictions.HostVector()[i],
+               precise_out_predictions.predictions.HostVector()[i]);
+    }
+
+    delete pp_precise;
+    delete pp_ellpack;
+  }
+
+  {
+    // Predictor should never try to create the histogram index by itself.  As only
+    // histogram index from training data is valid and predictor doesn't known which
+    // matrix is used for training.
+    auto pp_dmat = CreateDMatrix(rows, kCols, 0);
+    auto p_dmat = *pp_dmat;
+    PredictionCacheEntry precise_out_predictions;
+    predictor->PredictBatch(p_dmat.get(), &precise_out_predictions, model, 0);
+    ASSERT_FALSE(p_dmat->PageExists<Page>());
+    delete pp_dmat;
+  }
+}
+
+void TestTrainingPrediction(size_t rows, std::string tree_method);
+
+}  // namespace xgboost
+
+#endif  // XGBOOST_TEST_PREDICTOR_H_