Avoid thread block with sparse data. (#7255)

tests/cpp/predictor/test_cpu_predictor.cc

@@ -247,4 +247,9 @@ TEST(CpuPredictor, UpdatePredictionCache) {
 TEST(CpuPredictor, LesserFeatures) {
   TestPredictionWithLesserFeatures("cpu_predictor");
 }
+
+TEST(CpuPredictor, Sparse) {
+  TestSparsePrediction(0.2, "cpu_predictor");
+  TestSparsePrediction(0.8, "cpu_predictor");
+}
 }  // namespace xgboost

tests/cpp/predictor/test_gpu_predictor.cu

@@ -256,5 +256,10 @@ TEST(GPUPredictor, PredictLeafBasic) {
     ASSERT_EQ(v, 0);
   }
 }
+
+TEST(GPUPredictor, Sparse) {
+  TestSparsePrediction(0.2, "gpu_predictor");
+  TestSparsePrediction(0.8, "gpu_predictor");
+}
 }  // namespace predictor
 }  // namespace xgboost

tests/cpp/predictor/test_predictor.cc

@@ -11,6 +11,7 @@
 #include "test_predictor.h"
 
 #include "../helpers.h"
+#include "../../../src/data/adapter.h"
 #include "../../../src/common/io.h"
 #include "../../../src/common/categorical.h"
 #include "../../../src/common/bitfield.h"
@@ -355,4 +356,57 @@ void TestIterationRange(std::string name) {
     ASSERT_EQ(h_sliced, h_range);
   }
 }
+
+void TestSparsePrediction(float sparsity, std::string predictor) {
+  size_t constexpr kRows = 512, kCols = 128;
+  auto Xy = RandomDataGenerator(kRows, kCols, sparsity).GenerateDMatrix(true);
+  std::unique_ptr<Learner> learner{Learner::Create({Xy})};
+  learner->Configure();
+  for (size_t i = 0; i < 4; ++i) {
+    learner->UpdateOneIter(i, Xy);
+  }
+
+  HostDeviceVector<float> sparse_predt;
+
+  Json model{Object{}};
+  learner->SaveModel(&model);
+
+  learner.reset(Learner::Create({Xy}));
+  learner->LoadModel(model);
+
+  learner->SetParam("predictor", predictor);
+  learner->Predict(Xy, false, &sparse_predt, 0, 0);
+
+  std::vector<float> with_nan(kRows * kCols, std::numeric_limits<float>::quiet_NaN());
+  for (auto const& page : Xy->GetBatches<SparsePage>()) {
+    auto batch = page.GetView();
+    for (size_t i = 0; i < batch.Size(); ++i) {
+      auto row = batch[i];
+      for (auto e : row) {
+        with_nan[i * kCols + e.index] = e.fvalue;
+      }
+    }
+  }
+
+  learner->SetParam("predictor", "cpu_predictor");
+  // Xcode_12.4 doesn't compile with `std::make_shared`.
+  auto dense = std::shared_ptr<data::DenseAdapter>(
+      new data::DenseAdapter(with_nan.data(), kRows, kCols));
+  HostDeviceVector<float> *p_dense_predt;
+  learner->InplacePredict(dmlc::any(dense), nullptr, PredictionType::kValue,
+                          std::numeric_limits<float>::quiet_NaN(), &p_dense_predt,
+                          0, 0);
+
+  auto const& dense_predt = *p_dense_predt;
+  if (predictor == "cpu_predictor") {
+    ASSERT_EQ(dense_predt.HostVector(), sparse_predt.HostVector());
+  } else {
+    auto const &h_dense = dense_predt.HostVector();
+    auto const &h_sparse = sparse_predt.HostVector();
+    ASSERT_EQ(h_dense.size(), h_sparse.size());
+    for (size_t i = 0; i < h_dense.size(); ++i) {
+      ASSERT_FLOAT_EQ(h_dense[i], h_sparse[i]);
+    }
+  }
+}
 }  // namespace xgboost

tests/cpp/predictor/test_predictor.h

@@ -70,6 +70,8 @@ void TestCategoricalPrediction(std::string name);
 void TestCategoricalPredictLeaf(StringView name);
 
 void TestIterationRange(std::string name);
+
+void TestSparsePrediction(float sparsity, std::string predictor);
 }  // namespace xgboost
 
 #endif  // XGBOOST_TEST_PREDICTOR_H_