From 06d9b998ceca70d415b3d14f208332c28abdc02c Mon Sep 17 00:00:00 2001
From: amdsc21 <96135754+amdsc21@users.noreply.github.com>
Date: Tue, 28 Mar 2023 00:14:18 +0200
Subject: [PATCH] fix CAPI BuildInfo

---
 dmlc-core                                             | 2 +-
 src/c_api/c_api.cu                                    | 6 ++++++
 src/collective/communicator.cu                        | 4 ++--
 src/common/device_helpers.hip.h                       | 6 +++---
 tests/cpp/c_api/test_c_api.cc                         | 3 +++
 tests/cpp/collective/test_nccl_device_communicator.cu | 8 ++++++--
 6 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/dmlc-core b/dmlc-core
index dfd936526..ea21135fb 160000
--- a/dmlc-core
+++ b/dmlc-core
@@ -1 +1 @@
-Subproject commit dfd9365264a060a5096734b7d892e1858b6d2722
+Subproject commit ea21135fbb141ae103fb5fc960289b5601b468f2
diff --git a/src/c_api/c_api.cu b/src/c_api/c_api.cu
index 89830b89b..3fc772064 100644
--- a/src/c_api/c_api.cu
+++ b/src/c_api/c_api.cu
@@ -17,8 +17,11 @@ namespace xgboost {
 void XGBBuildInfoDevice(Json *p_info) {
   auto &info = *p_info;
 
+#if defined(XGBOOST_USE_CUDA)
   info["USE_CUDA"] = true;
+#elif defined(XGBOOST_USE_HIP)
   info["USE_HIP"] = true;
+#endif
 
   std::vector<Json> v{Json{Integer{THRUST_MAJOR_VERSION}}, Json{Integer{THRUST_MINOR_VERSION}},
                       Json{Integer{THRUST_SUBMINOR_VERSION}}};
@@ -29,6 +32,9 @@ void XGBBuildInfoDevice(Json *p_info) {
 
 #if defined(XGBOOST_USE_NCCL)
   info["USE_NCCL"] = Boolean{true};
+  v = {Json{Integer{NCCL_MAJOR}}, Json{Integer{NCCL_MINOR}}, Json{Integer{NCCL_PATCH}}};
+  info["NCCL_VERSION"] = v;
+#elif defined(XGBOOST_USE_RCCL)
   info["USE_RCCL"] = Boolean{true};
   v = {Json{Integer{NCCL_MAJOR}}, Json{Integer{NCCL_MINOR}}, Json{Integer{NCCL_PATCH}}};
   info["NCCL_VERSION"] = v;
diff --git a/src/collective/communicator.cu b/src/collective/communicator.cu
index 0880741f9..d0f6633c1 100644
--- a/src/collective/communicator.cu
+++ b/src/collective/communicator.cu
@@ -5,7 +5,7 @@
 #include "device_communicator.cuh"
 #include "device_communicator_adapter.cuh"
 #include "noop_communicator.h"
-#ifdef XGBOOST_USE_NCCL
+#if defined(XGBOOST_USE_NCCL) || defined(XGBOOST_USE_RCCL)
 #include "nccl_device_communicator.cuh"
 #endif
 
@@ -25,7 +25,7 @@ void Communicator::Finalize() {
 DeviceCommunicator* Communicator::GetDevice(int device_ordinal) {
   if (!device_communicator_ || device_ordinal_ != device_ordinal) {
     device_ordinal_ = device_ordinal;
-#ifdef XGBOOST_USE_NCCL
+#if defined(XGBOOST_USE_NCCL) || defined(XGBOOST_USE_RCCL)
     if (type_ != CommunicatorType::kFederated) {
       device_communicator_.reset(new NcclDeviceCommunicator(device_ordinal, Get()));
     } else {
diff --git a/src/common/device_helpers.hip.h b/src/common/device_helpers.hip.h
index 1d92bd332..b579bc5ac 100644
--- a/src/common/device_helpers.hip.h
+++ b/src/common/device_helpers.hip.h
@@ -38,9 +38,9 @@
 #include "xgboost/logging.h"
 #include "xgboost/span.h"
 
-#ifdef XGBOOST_USE_NCCL
-#include "nccl.h"
-#endif  // XGBOOST_USE_NCCL
+#ifdef XGBOOST_USE_RCCL
+#include "rccl.h"
+#endif  // XGBOOST_USE_RCCL
 
 #if defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
 #include "rmm/mr/device/per_device_resource.hpp"
diff --git a/tests/cpp/c_api/test_c_api.cc b/tests/cpp/c_api/test_c_api.cc
index a2595d360..b3dd9d3c4 100644
--- a/tests/cpp/c_api/test_c_api.cc
+++ b/tests/cpp/c_api/test_c_api.cc
@@ -362,10 +362,13 @@ TEST(CAPI, BuildInfo) {
   XGBuildInfo(&out);
   auto loaded = Json::Load(StringView{out});
   ASSERT_TRUE(get<Object const>(loaded).find("USE_OPENMP") != get<Object const>(loaded).cend());
+#if defined(XGBOOST_USE_CUDA)
   ASSERT_TRUE(get<Object const>(loaded).find("USE_CUDA") != get<Object const>(loaded).cend());
   ASSERT_TRUE(get<Object const>(loaded).find("USE_NCCL") != get<Object const>(loaded).cend());
+#elif defined(XGBOOST_USE_HIP)
   ASSERT_TRUE(get<Object const>(loaded).find("USE_HIP") != get<Object const>(loaded).cend());
   ASSERT_TRUE(get<Object const>(loaded).find("USE_RCCL") != get<Object const>(loaded).cend());
+#endif
 }
 
 TEST(CAPI, NullPtr) {
diff --git a/tests/cpp/collective/test_nccl_device_communicator.cu b/tests/cpp/collective/test_nccl_device_communicator.cu
index 8ce877aef..d75e020e3 100644
--- a/tests/cpp/collective/test_nccl_device_communicator.cu
+++ b/tests/cpp/collective/test_nccl_device_communicator.cu
@@ -1,13 +1,17 @@
 /**
  * Copyright 2022-2023, XGBoost contributors
  */
-#ifdef XGBOOST_USE_NCCL
+#if defined(XGBOOST_USE_NCCL) || defined(XGBOOST_USE_RCCL)
 
 #include <gtest/gtest.h>
 
 #include <string>  // for string
 
+#if defined(XGBOOST_USE_NCCL)
 #include "../../../src/collective/nccl_device_communicator.cuh"
+#elif defined(XGBOOST_USE_RCCL)
+#include "../../../src/collective/nccl_device_communicator.hip.h"
+#endif
 
 namespace xgboost {
 namespace collective {
@@ -33,4 +37,4 @@ TEST(NcclDeviceCommunicatorSimpleTest, SystemError) {
 }  // namespace collective
 }  // namespace xgboost
 
-#endif  // XGBOOST_USE_NCCL
+#endif  // XGBOOST_USE_NCCL || XGBOOST_USE_RCCL