deepmodeling · wanghan-iapcm · Oct 27, 2022 · Oct 11, 2022 · Oct 12, 2022 · Oct 12, 2022
diff --git a/.github/workflows/build_cc.yml b/.github/workflows/build_cc.yml
@@ -12,6 +12,7 @@ jobs:
         include:
         - variant: cpu
         - variant: cuda
+        - variant: rocm
         - variant: clang
     steps:
     - name: work around permission issue
@@ -23,6 +24,14 @@ jobs:
       if: matrix.variant == 'cpu'
     - run: apt-get update && apt-get install -y nvidia-cuda-toolkit
       if: matrix.variant == 'cuda'
+    - run: |
+         apt-get update && apt-get install -y gnupg2 \
+         && echo 'deb [arch=amd64] https://repo.radeon.com/rocm/apt/5.3/ jammy main' | tee /etc/apt/sources.list.d/rocm.list \
+         && printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | tee /etc/apt/preferences.d/rocm-pin-600 \
+         && curl -s https://repo.radeon.com/rocm/rocm.gpg.key | apt-key add - \
+         && apt-get update \
+         && apt-get install -y rocm-dev hipcub-dev
+      if: matrix.variant == 'rocm'
     - run: apt-get update && apt-get install -y clang
       if: matrix.variant == 'clang'  
     - run: source/install/build_cc.sh

diff --git a/doc/install/install-from-source.md b/doc/install/install-from-source.md
@@ -177,7 +177,7 @@ One may add the following arguments to `cmake`:
 | -DUSE_CUDA_TOOLKIT=&lt;value&gt; | `TRUE` or `FALSE` | `FALSE`       | If `TRUE`, Build GPU support with CUDA toolkit. |
 | -DCUDA_TOOLKIT_ROOT_DIR=&lt;value&gt; | Path         | Detected automatically | The path to the CUDA toolkit directory. CUDA 7.0 or later is supported. NVCC is required. |
 | -DUSE_ROCM_TOOLKIT=&lt;value&gt; | `TRUE` or `FALSE` | `FALSE`       | If `TRUE`, Build GPU support with ROCM toolkit. |
-| -DROCM_ROOT=&lt;value&gt; | Path         | Detected automatically | The path to the ROCM toolkit directory. |
+| -DCMAKE_HIP_COMPILER_ROCM_ROOT=&lt;value&gt; | Path         | Detected automatically | The path to the ROCM toolkit directory. |
 | -DLAMMPS_SOURCE_ROOT=&lt;value&gt; | Path         | - | Only neccessary for LAMMPS plugin mode. The path to the [LAMMPS source code](install-lammps.md). LAMMPS 8Apr2021 or later is supported. If not assigned, the plugin mode will not be enabled. |
 | -DUSE_TF_PYTHON_LIBS=&lt;value&gt; | `TRUE` or `FALSE` | `FALSE`       | If `TRUE`, Build C++ interface with TensorFlow's Python libraries(TensorFlow's Python Interface is required). And there's no need for building TensorFlow's C++ interface.|
 

diff --git a/setup.py b/setup.py
@@ -52,7 +52,7 @@
     cmake_args.append("-DUSE_ROCM_TOOLKIT:BOOL=TRUE")
     rocm_root = os.environ.get("ROCM_ROOT")
     if rocm_root:
-        cmake_args.append(f"-DROCM_ROOT:STRING={rocm_root}")
+        cmake_args.append(f"-DCMAKE_HIP_COMPILER_ROCM_ROOT:STRING={rocm_root}")
 else:
     raise RuntimeError("Unsupported DP_VARIANT option: %s" % dp_variant)
 

diff --git a/source/cmake/FindROCM.cmake b/source/cmake/FindROCM.cmake
@@ -7,13 +7,10 @@
 # ROCM_LIBRARIES
 
 # define the search path
-list(APPEND ROCM_search_PATHS ${ROCM_ROOT})
-list(APPEND ROCM_search_PATHS "/opt/rocm/")
-
-# define the libs to find
-if (NOT ROCM_FIND_COMPONENTS)
-  set(ROCM_FIND_COMPONENTS hip_hcc hiprtc)
-endif ()
+cmake_minimum_required(VERSION 3.21)
+include(CMakeDetermineHIPCompiler)
+set(ROCM_PATH ${CMAKE_HIP_COMPILER_ROCM_ROOT})
+set(ROCM_search_PATHS ${CMAKE_HIP_COMPILER_ROCM_ROOT})
 
 # includes
 find_path (ROCM_INCLUDE_DIRS
@@ -31,21 +28,6 @@ if (NOT ROCM_INCLUDE_DIRS AND ROCM_FIND_REQUIRED)
     "You can manually set the ROCM install path by -DROCM_ROOT ")
 endif ()
 
-# libs
-foreach (module ${ROCM_FIND_COMPONENTS})
-  find_library(ROCM_LIBRARIES_${module}
-    NAMES ${module}
-    PATHS ${ROCM_search_PATHS} PATH_SUFFIXES "lib" NO_DEFAULT_PATH
-    )
-  if (ROCM_LIBRARIES_${module})
-    list(APPEND ROCM_LIBRARIES ${ROCM_LIBRARIES_${module}})
-  elseif (ROCM_FIND_REQUIRED)
-    message(FATAL_ERROR 
-      "Not found lib/'${module}' in '${ROCM_search_PATHS}' "
-      "You can manually set the ROCM install path by -DROCM_ROOT ")
-  endif ()
-endforeach ()
-
 # FindHIP.cmake
 find_path (HIP_CMAKE
   NAMES 
@@ -64,6 +46,30 @@ endif ()
 list (APPEND CMAKE_MODULE_PATH ${HIP_CMAKE})
 find_package(HIP) 
 
+# define the libs to find
+if (NOT ROCM_FIND_COMPONENTS)
+  if (HIP_VERSION VERSION_GREATER_EQUAL 3.5.1)
+    set(ROCM_FIND_COMPONENTS amd_comgr amdhip64)
+  else()
+    set(ROCM_FIND_COMPONENTS hip-hcc hiprtc)
+  endif()
+endif ()
+
+# libs
+foreach (module ${ROCM_FIND_COMPONENTS})
+  find_library(ROCM_LIBRARIES_${module}
+    NAMES ${module}
+    PATHS ${ROCM_search_PATHS} PATH_SUFFIXES "lib" NO_DEFAULT_PATH
+    )
+  if (ROCM_LIBRARIES_${module})
+    list(APPEND ROCM_LIBRARIES ${ROCM_LIBRARIES_${module}})
+  elseif (ROCM_FIND_REQUIRED)
+    message(FATAL_ERROR 
+      "Not found lib/'${module}' in '${ROCM_search_PATHS}' "
+      "You can manually set the ROCM install path by -DROCM_ROOT ")
+  endif ()
+endforeach ()
+
 # define the output variable
 if (ROCM_INCLUDE_DIRS AND ROCM_LIBRARIES AND HIP_CMAKE)
   set(ROCM_FOUND TRUE)

diff --git a/source/install/build_cc.sh b/source/install/build_cc.sh
@@ -1,8 +1,11 @@
 set -e
 
-if [ "$DP_VARIANT" == "cuda" ]
+if [ "$DP_VARIANT" = "cuda" ]
 then
   CUDA_ARGS="-DUSE_CUDA_TOOLKIT=TRUE"
+elif [ "$DP_VARIANT" = "rocm" ]
+then
+  CUDA_ARGS="-DUSE_ROCM_TOOLKIT=TRUE"
 fi
 #------------------
 

diff --git a/source/lib/src/rocm/CMakeLists.txt b/source/lib/src/rocm/CMakeLists.txt
@@ -1,18 +1,21 @@
 # required cmake version
-cmake_minimum_required(VERSION 3.5)
+cmake_minimum_required(VERSION 3.21)
 # project name
 project(deepmd_op_rocm)
 set(CMAKE_LINK_WHAT_YOU_USE TRUE)
 
 # set c++ version c++11
-SET(CMAKE_CXX_STANDARD 11)
-SET(CMAKE_HIP_STANDARD 11)
+set(CMAKE_CXX_STANDARD 14)
+set(CMAKE_HIP_STANDARD 14)
 add_definitions("-DCUB_IGNORE_DEPRECATED_CPP_DIALECT")
 add_definitions("-DCUB_IGNORE_DEPRECATED_CPP_DIALECT")
 
 message(STATUS "HIP major version is " ${HIP_VERSION_MAJOR})
 
-set (HIP_HIPCC_FLAGS -hc; -fno-gpu-rdc; --amdgpu-target=gfx906; -fPIC; -O3; --std=c++11)
+set (HIP_HIPCC_FLAGS -fno-gpu-rdc; -fPIC --std=c++14) # --amdgpu-target=gfx906
+if (HIP_VERSION VERSION_LESS 3.5.1)
+  set (HIP_HIPCC_FLAGS -hc; ${HIP_HIPCC_FLAGS})
+endif()
 
 file (GLOB SOURCE_FILES "*.hip.cu" )
 

diff --git a/source/lib/src/rocm/neighbor_list.hip.cu b/source/lib/src/rocm/neighbor_list.hip.cu
@@ -288,7 +288,7 @@ void use_nei_info_gpu_rocm(
     dim3 block_grid(nloc, nblock);
     dim3 thread_grid(1, TPB);
     DPErrcheck(hipMemset(ntype, 0, sizeof(int) * nloc * nnei));
-    DPErrcheck(hipMemset(nmask, 0, sizeof(FPTYPE) * nloc * nnei));
+    DPErrcheck(hipMemset(nmask, 0, sizeof(bool) * nloc * nnei));
     if (b_nlist_map){
         hipLaunchKernelGGL(map_nei_info, block_grid, thread_grid, 0, 0, nlist, ntype, nmask, type, nlist_map, nloc, nnei, ntypes);
     }
@@ -301,4 +301,4 @@ void use_nei_info_gpu_rocm(
 
 template int build_nlist_gpu_rocm<float>(InputNlist & nlist, int * max_list_size, int * nlist_data, const float * c_cpy, const int & nloc, const int & nall, const int & mem_size, const float & rcut);
 template int build_nlist_gpu_rocm<double>(InputNlist & nlist, int * max_list_size, int * nlist_data, const double * c_cpy, const int & nloc, const int & nall, const int & mem_size, const float & rcut);
-}
+}