Fix specifying gpu_id, add tests. (#3851)

* Rewrite gpu_id related code.

* Remove normalised/unnormalised operatios.
* Address difference between `Index' and `Device ID'.
* Modify doc for `gpu_id'.
* Better LOG for GPUSet.
* Check specified n_gpus.
* Remove inappropriate `device_idx' term.
* Clarify GpuIdType and size_t.

tests/cpp/common/test_common.cc

@@ -10,11 +10,7 @@ TEST(GPUSet, Basic) {
   ASSERT_TRUE(devices != GPUSet::Empty());
   EXPECT_EQ(devices.Size(), 1);
 
-  EXPECT_ANY_THROW(devices.Index(1));
-  EXPECT_ANY_THROW(devices.Index(-1));
-
   devices = GPUSet::Range(1, 0);
-  EXPECT_EQ(devices, GPUSet::Empty());
   EXPECT_EQ(devices.Size(), 0);
   EXPECT_TRUE(devices.IsEmpty());
 
@@ -25,18 +21,17 @@ TEST(GPUSet, Basic) {
   EXPECT_EQ(devices.Size(), 0);
   EXPECT_TRUE(devices.IsEmpty());
 
-  devices = GPUSet::Range(2, 8);
+  devices = GPUSet::Range(2, 8);  // 2 ~ 10
   EXPECT_EQ(devices.Size(), 8);
-  EXPECT_ANY_THROW(devices[8]);
-  EXPECT_ANY_THROW(devices.Index(0));
+  EXPECT_ANY_THROW(devices.DeviceId(8));
 
-  devices = devices.Unnormalised();
+  auto device_id = devices.DeviceId(0);
+  EXPECT_EQ(device_id, 2);
+  auto device_index = devices.Index(2);
+  EXPECT_EQ(device_index, 0);
 
-  EXPECT_EQ(*devices.begin(), 0);
-  EXPECT_EQ(*devices.end(), devices.Size());
 #ifndef XGBOOST_USE_CUDA
   EXPECT_EQ(GPUSet::AllVisible(), GPUSet::Empty());
 #endif
 }
 }  // namespace xgboost
-

tests/cpp/common/test_common.cu

@@ -7,12 +7,10 @@ TEST(GPUSet, GPUBasic) {
   GPUSet devices = GPUSet::Empty();
   ASSERT_TRUE(devices.IsEmpty());
 
-  devices = GPUSet{0, 1};
+  devices = GPUSet{1, 1};
   ASSERT_TRUE(devices != GPUSet::Empty());
   EXPECT_EQ(devices.Size(), 1);
-
-  EXPECT_ANY_THROW(devices.Index(1));
-  EXPECT_ANY_THROW(devices.Index(-1));
+  EXPECT_EQ(*(devices.begin()), 1);
 
   devices = GPUSet::Range(1, 0);
   EXPECT_EQ(devices, GPUSet::Empty());
@@ -23,15 +21,12 @@ TEST(GPUSet, GPUBasic) {
 
   devices = GPUSet::Range(2, -1);
   EXPECT_EQ(devices, GPUSet::Empty());
-  EXPECT_EQ(devices.Size(), 0);
-  EXPECT_TRUE(devices.IsEmpty());
 
   devices = GPUSet::Range(2, 8);
   EXPECT_EQ(devices.Size(), 8);
-  devices = devices.Unnormalised();
 
-  EXPECT_EQ(*devices.begin(), 0);
-  EXPECT_EQ(*devices.end(), devices.Size());
+  EXPECT_EQ(*devices.begin(), 2);
+  EXPECT_EQ(*devices.end(), 2 + devices.Size());
   EXPECT_EQ(8, devices.Size());
 
   ASSERT_NO_THROW(GPUSet::AllVisible());
@@ -41,4 +36,27 @@ TEST(GPUSet, GPUBasic) {
   }
 }
 
-}  // namespace xgboost
+#if defined(XGBOOST_USE_NCCL)
+TEST(GPUSet, MGPU_GPUBasic) {
+  {
+    GPUSet devices = GPUSet::All(1, 1);
+    ASSERT_EQ(*(devices.begin()), 1);
+    ASSERT_EQ(*(devices.end()), 2);
+    ASSERT_EQ(devices.Size(), 1);
+    ASSERT_TRUE(devices.Contains(1));
+  }
+
+  {
+    GPUSet devices = GPUSet::All(0, -1);
+    ASSERT_GE(devices.Size(), 2);
+  }
+
+  // Specify number of rows.
+  {
+    GPUSet devices = GPUSet::All(0, -1, 1);
+    ASSERT_EQ(devices.Size(), 1);
+  }
+}
+#endif
+
+}  // namespace xgboost

tests/cpp/common/test_transform_range.cu

@@ -38,6 +38,73 @@ TEST(Transform, MGPU_Basic) {
   ASSERT_TRUE(std::equal(h_sol.begin(), h_sol.end(), res.begin()));
 }
 
+// Test for multi-classes setting.
+template <typename T>
+struct TestTransformRangeGranular {
+  const size_t granularity = 8;
+
+  TestTransformRangeGranular(const size_t granular) : granularity{granular} {}
+  void XGBOOST_DEVICE operator()(size_t _idx,
+                                 Span<bst_float> _out, Span<const bst_float> _in) {
+    auto in_sub = _in.subspan(_idx * granularity, granularity);
+    auto out_sub = _out.subspan(_idx * granularity, granularity);
+    for (size_t i = 0; i < granularity; ++i) {
+      out_sub[i] = in_sub[i];
+    }
+  }
+};
+
+TEST(Transform, MGPU_Granularity) {
+  GPUSet devices = GPUSet::All(0, -1);
+
+  const size_t size {8990};
+  const size_t granularity = 10;
+
+  GPUDistribution distribution =
+      GPUDistribution::Granular(devices, granularity);
+
+  std::vector<bst_float> h_in(size);
+  std::vector<bst_float> h_out(size);
+  InitializeRange(h_in.begin(), h_in.end());
+  std::vector<bst_float> h_sol(size);
+  InitializeRange(h_sol.begin(), h_sol.end());
+
+  const HostDeviceVector<bst_float> in_vec {h_in, distribution};
+  HostDeviceVector<bst_float> out_vec {h_out, distribution};
+
+  ASSERT_NO_THROW(
+      Transform<>::Init(
+          TestTransformRangeGranular<bst_float>{granularity},
+          Range{0, size / granularity},
+          distribution)
+      .Eval(&out_vec, &in_vec));
+  std::vector<bst_float> res = out_vec.HostVector();
+
+  ASSERT_TRUE(std::equal(h_sol.begin(), h_sol.end(), res.begin()));
+}
+
+TEST(Transform, MGPU_SpecifiedGpuId) {
+  // Use 1 GPU, Numbering of GPU starts from 1
+  auto devices = GPUSet::All(1, 1);
+  const size_t size {256};
+  std::vector<bst_float> h_in(size);
+  std::vector<bst_float> h_out(size);
+  InitializeRange(h_in.begin(), h_in.end());
+  std::vector<bst_float> h_sol(size);
+  InitializeRange(h_sol.begin(), h_sol.end());
+
+  const HostDeviceVector<bst_float> in_vec {h_in,
+        GPUDistribution::Block(devices)};
+  HostDeviceVector<bst_float> out_vec {h_out,
+        GPUDistribution::Block(devices)};
+
+  ASSERT_NO_THROW(
+      Transform<>::Init(TestTransformRange<bst_float>{}, Range{0, size}, devices)
+      .Eval(&out_vec, &in_vec));
+  std::vector<bst_float> res = out_vec.HostVector();
+  ASSERT_TRUE(std::equal(h_sol.begin(), h_sol.end(), res.begin()));
+}
+
 }  // namespace xgboost
 }  // namespace common
 #endif