Refactor device communicator to make allreduce more flexible (#9295)

2023-06-13 12:53:03 -07:00
parent c2f0486d37
commit e70810be8a
11 changed files with 190 additions and 106 deletions
--- a/tests/cpp/plugin/test_federated_adapter.cu
+++ b/tests/cpp/plugin/test_federated_adapter.cu
@@ -36,7 +36,7 @@ TEST_F(FederatedAdapterTest, DeviceAllReduceSum) {
      int count = 3;
      thrust::device_vector<double> buffer(count, 0);
      thrust::sequence(buffer.begin(), buffer.end());
-      adapter.AllReduceSum(buffer.data().get(), count);
+      adapter.AllReduce(buffer.data().get(), count, DataType::kDouble, Operation::kSum);
      thrust::host_vector<double> host_buffer = buffer;
      EXPECT_EQ(host_buffer.size(), count);
      for (auto i = 0; i < count; i++) {