From e0716afabfb322f908045367d8a683bffb9f5f9a Mon Sep 17 00:00:00 2001
From: amdsc21 <96135754+amdsc21@users.noreply.github.com>
Date: Thu, 23 Mar 2023 20:22:34 +0100
Subject: [PATCH] fix objective/objective.cc, CMakeFile and setup.py

---
 CMakeLists.txt             | 5 +++++
 python-package/setup.py    | 8 ++++++++
 rocgputreeshap             | 2 +-
 src/objective/objective.cc | 2 +-
 warp-primitives            | 2 +-
 5 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index e6a3c4bd4..4cc47fa6a 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -187,6 +187,7 @@ if (USE_HIP)
   set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -Wunused-result -w")
   set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} -D__HIP_PLATFORM_AMD__")
   add_subdirectory(${PROJECT_SOURCE_DIR}/rocgputreeshap)
+  add_subdirectory(${PROJECT_SOURCE_DIR}/warp-primitives)
 endif (USE_HIP)
 
 if (FORCE_COLORED_OUTPUT AND (CMAKE_GENERATOR STREQUAL "Ninja") AND
@@ -228,6 +229,10 @@ if (USE_NCCL)
   find_package(Nccl REQUIRED)
 endif (USE_NCCL)
 
+if (USE_RCCL)
+  find_package(rccl REQUIRED)
+endif (USE_RCCL)
+
 # dmlc-core
 msvc_use_static_runtime()
 if (FORCE_SHARED_CRT)
diff --git a/python-package/setup.py b/python-package/setup.py
index fe1cbf2e9..006a2ea69 100644
--- a/python-package/setup.py
+++ b/python-package/setup.py
@@ -23,6 +23,8 @@ USER_OPTIONS = {
     "use-cuda": (None, "Build with GPU acceleration.", 0),
     "use-nccl": (None, "Build with NCCL to enable distributed GPU support.", 0),
     "build-with-shared-nccl": (None, "Build with shared NCCL library.", 0),
+    "use-hip": (None, "Build with GPU acceleration.", 0),
+    "use-rccl": (None, "Build with RCCL to enable distributed GPU support.", 0),
     "hide-cxx-symbols": (None, "Hide all C++ symbols during build.", 1),
     "use-hdfs": (None, "Build with HDFS support", 0),
     "use-azure": (None, "Build with AZURE support.", 0),
@@ -65,6 +67,8 @@ def copy_tree(src_dir: str, target_dir: str) -> None:
     inc = os.path.join(src_dir, "include")
     dmlc_core = os.path.join(src_dir, "dmlc-core")
     gputreeshap = os.path.join(src_dir, "gputreeshap")
+    rocgputreeshap = os.path.join(src_dir, "rocgputreeshap")
+    warpprim= os.path.join(src_dir, "warp-primitives")
     rabit = os.path.join(src_dir, "rabit")
     cmake = os.path.join(src_dir, "cmake")
     plugin = os.path.join(src_dir, "plugin")
@@ -73,6 +77,8 @@ def copy_tree(src_dir: str, target_dir: str) -> None:
     clean_copy_tree(inc, os.path.join(target_dir, "include"))
     clean_copy_tree(dmlc_core, os.path.join(target_dir, "dmlc-core"))
     clean_copy_tree(gputreeshap, os.path.join(target_dir, "gputreeshap"))
+    clean_copy_tree(rocgputreeshap, os.path.join(target_dir, "rocgputreeshap"))
+    clean_copy_tree(warpprim, os.path.join(target_dir, "warp-primitives"))
     clean_copy_tree(rabit, os.path.join(target_dir, "rabit"))
     clean_copy_tree(cmake, os.path.join(target_dir, "cmake"))
     clean_copy_tree(plugin, os.path.join(target_dir, "plugin"))
@@ -306,6 +312,8 @@ class Install(install.install):  # pylint: disable=too-many-instance-attributes
         self.use_cuda = 0
         self.use_nccl = 0
         self.build_with_shared_nccl = 0
+        self.use_hip= 0
+        self.use_rccl = 0
         self.hide_cxx_symbols = 1
 
         self.use_hdfs = 0
diff --git a/rocgputreeshap b/rocgputreeshap
index 0ce793d34..3704f6142 160000
--- a/rocgputreeshap
+++ b/rocgputreeshap
@@ -1 +1 @@
-Subproject commit 0ce793d3476d3d1a36256a6beb40626748cac608
+Subproject commit 3704f6142138766bb6e3585f496c8b7de61d2d32
diff --git a/src/objective/objective.cc b/src/objective/objective.cc
index 70746a1f3..925456fd0 100644
--- a/src/objective/objective.cc
+++ b/src/objective/objective.cc
@@ -42,7 +42,7 @@ void ObjFunction::InitEstimation(MetaInfo const&, linalg::Tensor<float, 1>* base
 namespace xgboost {
 namespace obj {
 // List of files that will be force linked in static links.
-#if !defined(XGBOOST_USE_CUDA) && !defined(XGBOOST_USE_HIP)
+#if defined(XGBOOST_USE_CUDA) || defined(XGBOOST_USE_HIP)
 DMLC_REGISTRY_LINK_TAG(regression_obj_gpu);
 DMLC_REGISTRY_LINK_TAG(quantile_obj_gpu);
 DMLC_REGISTRY_LINK_TAG(hinge_obj_gpu);
diff --git a/warp-primitives b/warp-primitives
index d8d1bb6ff..af1eccf83 160000
--- a/warp-primitives
+++ b/warp-primitives
@@ -1 +1 @@
-Subproject commit d8d1bb6fff784e3c30f42d22d1fe09ca18c4c2e7
+Subproject commit af1eccf8313f0579ff190d4b76627b4559f19d1a