From 5037abeb86ceb12e008cfba44ebe980b72d893f6 Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Mon, 19 Oct 2020 12:02:36 +0800
Subject: [PATCH] Fix linear gpu input (#6255)

---
 src/gbm/gblinear.cc                 |  5 +++--
 tests/cpp/gbm/test_gblinear.cc      |  1 -
 tests/python-gpu/test_gpu_linear.py | 21 +++++++++++++++++++++
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/gbm/gblinear.cc b/src/gbm/gblinear.cc
index b23cdb449..728265b81 100644
--- a/src/gbm/gblinear.cc
+++ b/src/gbm/gblinear.cc
@@ -231,7 +231,8 @@ class GBLinear : public GradientBooster {
     // start collecting the prediction
     const int ngroup = model_.learner_model_param->num_output_group;
     preds.resize(p_fmat->Info().num_row_ * ngroup);
-    for (const auto &batch : p_fmat->GetBatches<SparsePage>()) {
+    for (const auto &page : p_fmat->GetBatches<SparsePage>()) {
+      auto const& batch = page.GetView();
       // output convention: nrow * k, where nrow is number of rows
       // k is number of group
       // parallel over local batch
@@ -241,7 +242,7 @@ class GBLinear : public GradientBooster {
       }
 #pragma omp parallel for schedule(static)
       for (omp_ulong i = 0; i < nsize; ++i) {
-        const size_t ridx = batch.base_rowid + i;
+        const size_t ridx = page.base_rowid + i;
         // loop over output groups
         for (int gid = 0; gid < ngroup; ++gid) {
           bst_float margin =
diff --git a/tests/cpp/gbm/test_gblinear.cc b/tests/cpp/gbm/test_gblinear.cc
index 481e4a726..61d22f5ea 100644
--- a/tests/cpp/gbm/test_gblinear.cc
+++ b/tests/cpp/gbm/test_gblinear.cc
@@ -44,6 +44,5 @@ TEST(GBLinear, JsonIO) {
     ASSERT_EQ(weights.size(), 17);
   }
 }
-
 }  // namespace gbm
 }  // namespace xgboost
diff --git a/tests/python-gpu/test_gpu_linear.py b/tests/python-gpu/test_gpu_linear.py
index b887175df..262a4eb95 100644
--- a/tests/python-gpu/test_gpu_linear.py
+++ b/tests/python-gpu/test_gpu_linear.py
@@ -1,5 +1,7 @@
 import sys
 from hypothesis import strategies, given, settings, assume
+import pytest
+import numpy
 import xgboost as xgb
 sys.path.append("tests/python")
 import testing as tm
@@ -48,3 +50,22 @@ class TestGPULinear:
         param = dataset.set_params(param)
         result = train_result(param, dataset.get_dmat(), num_rounds)['train'][dataset.metric]
         assert tm.non_increasing([result[0], result[-1]])
+
+    @pytest.mark.skipif(**tm.no_cupy())
+    def test_gpu_coordinate_from_cupy(self):
+        # Training linear model is quite expensive, so we don't include it in
+        # test_from_cupy.py
+        import cupy
+        params = {'booster': 'gblinear', 'updater': 'gpu_coord_descent',
+                  'n_estimators': 100}
+        X, y = tm.get_boston()
+        cpu_model = xgb.XGBRegressor(**params)
+        cpu_model.fit(X, y)
+        cpu_predt = cpu_model.predict(X)
+
+        X = cupy.array(X)
+        y = cupy.array(y)
+        gpu_model = xgb.XGBRegressor(**params)
+        gpu_model.fit(X, y)
+        gpu_predt = gpu_model.predict(X)
+        cupy.testing.assert_allclose(cpu_predt, gpu_predt)