From 8d160a206e2ca8accc31ceb65f2cdd7315d23711 Mon Sep 17 00:00:00 2001
From: Hui Liu <96135754+amdsc21@users.noreply.github.com>
Date: Mon, 30 Oct 2023 11:49:47 -0700
Subject: [PATCH] add jvm rocm support

---
 jvm-packages/CMakeLists.txt                          |  5 +++++
 jvm-packages/create_jni.py                           | 12 +++++++++---
 jvm-packages/pom.xml                                 |  1 +
 jvm-packages/xgboost4j-gpu/pom.xml                   |  2 ++
 .../xgboost4j-gpu/src/native/xgboost4j-gpu.cpp       |  2 +-
 .../xgboost4j-gpu/src/native/xgboost4j-gpu.cu        |  4 ++++
 .../xgboost4j-gpu/src/native/xgboost4j-gpu.hip       |  4 ++++
 7 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/jvm-packages/CMakeLists.txt b/jvm-packages/CMakeLists.txt
index e1d0e94e2..f9b964a89 100644
--- a/jvm-packages/CMakeLists.txt
+++ b/jvm-packages/CMakeLists.txt
@@ -9,6 +9,11 @@ if (USE_CUDA)
     ${PROJECT_SOURCE_DIR}/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cu)
 endif (USE_CUDA)
 
+if (USE_HIP)
+  list(APPEND JVM_SOURCES
+    ${PROJECT_SOURCE_DIR}/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.hip)
+endif (USE_HIP)
+
 add_library(xgboost4j SHARED ${JVM_SOURCES} ${XGBOOST_OBJ_SOURCES})
 
 if (ENABLE_ALL_WARNINGS)
diff --git a/jvm-packages/create_jni.py b/jvm-packages/create_jni.py
index 18908fc1c..009d0cf6d 100755
--- a/jvm-packages/create_jni.py
+++ b/jvm-packages/create_jni.py
@@ -22,6 +22,8 @@ CONFIG = {
 
     "USE_CUDA": "OFF",
     "USE_NCCL": "OFF",
+    "USE_HIP": "OFF",
+    "USE_RCCL": "OFF",
     "JVM_BINDINGS": "ON",
     "LOG_CAPI_INVOCATION": "OFF"
 }
@@ -74,6 +76,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('--log-capi-invocation', type=str, choices=['ON', 'OFF'], default='OFF')
     parser.add_argument('--use-cuda', type=str, choices=['ON', 'OFF'], default='OFF')
+    parser.add_argument('--use-hip', type=str, choices=['ON', 'OFF'], default='OFF')
     cli_args = parser.parse_args()
 
     if sys.platform == "darwin":
@@ -84,7 +87,7 @@ if __name__ == "__main__":
 
     print("building Java wrapper")
     with cd(".."):
-        build_dir = 'build-gpu' if cli_args.use_cuda == 'ON' else 'build'
+        build_dir = 'build-gpu' if cli_args.use_cuda == 'ON' or cli_args.use_hip == 'ON' else 'build'
         maybe_makedirs(build_dir)
         with cd(build_dir):
             if sys.platform == "win32":
@@ -103,6 +106,9 @@ if __name__ == "__main__":
             if cli_args.use_cuda == 'ON':
                 CONFIG['USE_CUDA'] = 'ON'
                 CONFIG['USE_NCCL'] = 'ON'
+            elif cli_args.use_hip== 'ON':
+                CONFIG['USE_HIP'] = 'ON'
+                CONFIG['USE_RCCL'] = 'ON'
 
             args = ["-D{0}:BOOL={1}".format(k, v) for k, v in CONFIG.items()]
 
@@ -125,8 +131,8 @@ if __name__ == "__main__":
             run(f'"{sys.executable}" mapfeat.py')
             run(f'"{sys.executable}" mknfold.py machine.txt 1')
 
-    xgboost4j = 'xgboost4j-gpu' if cli_args.use_cuda == 'ON' else 'xgboost4j'
-    xgboost4j_spark = 'xgboost4j-spark-gpu' if cli_args.use_cuda == 'ON' else 'xgboost4j-spark'
+    xgboost4j = 'xgboost4j-gpu' if cli_args.use_cuda == 'ON' or cli_args.use_hip== 'ON' else 'xgboost4j'
+    xgboost4j_spark = 'xgboost4j-spark-gpu' if cli_args.use_cuda == 'ON' or cli_args.use_hip == 'ON' else 'xgboost4j-spark'
 
     print("copying native library")
     library_name, os_folder = {
diff --git a/jvm-packages/pom.xml b/jvm-packages/pom.xml
index 0faf52b8e..ec73605b4 100644
--- a/jvm-packages/pom.xml
+++ b/jvm-packages/pom.xml
@@ -43,6 +43,7 @@
         <maven.wagon.http.retryHandler.count>5</maven.wagon.http.retryHandler.count>
         <log.capi.invocation>OFF</log.capi.invocation>
         <use.cuda>OFF</use.cuda>
+        <use.hip>OFF</use.hip>
         <cudf.version>23.08.0</cudf.version>
         <spark.rapids.version>23.08.0</spark.rapids.version>
         <cudf.classifier>cuda11</cudf.classifier>
diff --git a/jvm-packages/xgboost4j-gpu/pom.xml b/jvm-packages/xgboost4j-gpu/pom.xml
index b9ff1590c..e46df9afb 100644
--- a/jvm-packages/xgboost4j-gpu/pom.xml
+++ b/jvm-packages/xgboost4j-gpu/pom.xml
@@ -104,6 +104,8 @@
                                 <argument>${log.capi.invocation}</argument>
                                 <argument>--use-cuda</argument>
                                 <argument>${use.cuda}</argument>
+                                <argument>--use-hip</argument>
+                                <argument>${use.hip}</argument>
                             </arguments>
                             <workingDirectory>${user.dir}</workingDirectory>
                         </configuration>
diff --git a/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cpp b/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cpp
index 698da6244..57769e5dc 100644
--- a/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cpp
+++ b/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cpp
@@ -2,7 +2,7 @@
 // Created by bobwang on 2021/9/8.
 //
 
-#ifndef XGBOOST_USE_CUDA
+#if !defined(XGBOOST_USE_CUDA) && !defined(XGBOOST_USE_HIP)
 
 #include <jni.h>
 
diff --git a/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cu b/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cu
index 317be01ad..272a90354 100644
--- a/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cu
+++ b/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.cu
@@ -1,6 +1,10 @@
 #include <jni.h>
 
+#if defined(XGBOOST_USE_CUDA)
 #include "../../../../src/common/device_helpers.cuh"
+#elif defined(XGBOOST_USE_HIP)
+#include "../../../../src/common/device_helpers.hip.h"
+#endif
 #include "../../../../src/common/cuda_pinned_allocator.h"
 #include "../../../../src/data/array_interface.h"
 #include "jvm_utils.h"
diff --git a/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.hip b/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.hip
index e69de29bb..2095d4182 100644
--- a/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.hip
+++ b/jvm-packages/xgboost4j-gpu/src/native/xgboost4j-gpu.hip
@@ -0,0 +1,4 @@
+
+#if defined(XGBOOST_USE_HIP)
+#include "xgboost4j-gpu.cu"
+#endif