Add ROCM support for DEEPMD

source/CMakeLists.txt

@@ -51,26 +51,30 @@ file(READ ${PROJECT_SOURCE_DIR}/config/MODEL_VER MODEL_VERSION)
 string(REPLACE "\n" " " MODEL_VERSION ${MODEL_VERSION})
 message(STATUS "Supported model version: ${MODEL_VERSION}")
 
-# define USE_CUDA_TOOLKIT
-if (DEFINED USE_CUDA_TOOLKIT)
-  if (USE_CUDA_TOOLKIT)
-    find_package(CUDA REQUIRED)
-  else()
-    message(STATUS "Will not build nv GPU support")
-  endif()
-else()
-  find_package(CUDA QUIET)  
-  if (CUDA_FOUND)
-    set(USE_CUDA_TOOLKIT TRUE)
-    message(STATUS "Found CUDA in ${CUDA_TOOLKIT_ROOT_DIR}, build nv GPU support")
-  else()
-    set(USE_CUDA_TOOLKIT FALSE)
-    message(STATUS "No cuda support found, will not build nv GPU support")
-  endif()
+# Devices that have both ROCM and CUDA are not currently supported
+if (USE_ROCM_TOOLKIT AND USE_CUDA_TOOLKIT)
+  message (FATAL_ERROR "Devices that have both ROCM and CUDA are not currently supported")
 endif()
+
+# define USE_CUDA_TOOLKIT
 if (USE_CUDA_TOOLKIT)
-  add_definitions("-D GOOGLE_CUDA")
-endif()
+  find_package(CUDA REQUIRED)
+  add_definitions("-DGOOGLE_CUDA")
+  message(STATUS "Found CUDA in ${CUDA_TOOLKIT_ROOT_DIR}, build nv GPU support")
+else()
+  message(STATUS "Will not build nv GPU support")
+endif(USE_CUDA_TOOLKIT)
+
+#define USE_ROCM_TOOLKIT
+if (USE_ROCM_TOOLKIT)
+  find_package(ROCM REQUIRED)
+  add_definitions("-DTENSORFLOW_USE_ROCM")
+  add_compile_definitions(__HIP_PLATFORM_HCC__)
+  message(STATUS "Found ROCM in ${ROCM_ROOT}, build AMD GPU support")
+else()
+  message(STATUS "Will not build AMD GPU support")  
+endif (USE_ROCM_TOOLKIT)
+
 
 # find tensorflow, I need tf abi info
 find_package(tensorflow REQUIRED)
@@ -195,9 +199,11 @@ if (BUILD_CPP_IF)
   set (LIB_DEEPMD_OP		"deepmd_op")
   set (LIB_DEEPMD_CC		"deepmd_cc")
   if (USE_CUDA_TOOLKIT)
-    set (LIB_DEEPMD_OP_CUDA		"deepmd_op_cuda")
-  else ()
-    set (LIB_DEEPMD_OP_CUDA		"deepmd_op")
+    set (LIB_DEEPMD_OP_DEVICE		"deepmd_op_cuda")
+  elseif (USE_ROCM_TOOLKIT)
+    set (LIB_DEEPMD_OP_DEVICE   "deepmd_op_rocm")
+  else()
+    set (LIB_DEEPMD_OP_DEVICE		"deepmd_op")
   endif()
   if (CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 4.9)
     set (LIB_DEEPMD_NATIVE	"deepmd_native_md")

source/api_cc/CMakeLists.txt

@@ -1,5 +1,4 @@
 # libmd
-
 set (libname ${LIB_DEEPMD_CC})
 
 configure_file(
@@ -12,6 +11,10 @@ if (USE_CUDA_TOOLKIT)
   include_directories("${CUDA_INCLUDE_DIRS}")
 endif()
 
+if (USE_ROCM_TOOLKIT)
+  include_directories("${ROCM_INCLUDE_DIRS}")
+endif()
+
 file(GLOB LIB_SRC src/*.cc src/*.cpp)
 file(GLOB INC_SRC include/*.h ${CMAKE_CURRENT_BINARY_DIR}/version.h)
 
@@ -21,6 +24,11 @@ if (USE_CUDA_TOOLKIT)
   target_link_libraries (${libname} ${CUDA_LIBRARIES})
 endif()
 
+if (USE_ROCM_TOOLKIT)
+  target_link_libraries (${libname} ${ROCM_LIBRARIES})
+endif()
+
+
 install(TARGETS ${libname} DESTINATION lib/)
 
 install(

source/api_cc/include/custom_op.h

@@ -23,12 +23,12 @@ struct DeviceFunctor {
   {
     device = "CPU";
   }
-  #if GOOGLE_CUDA
+  #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
   void operator()(
       std::string& device, 
       const GPUDevice& d) 
   {
     device = "GPU";
   }
-  #endif // GOOGLE_CUDA
+  #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 };

source/api_cc/src/DeepPot.cc

@@ -19,6 +19,19 @@ inline void cudaAssert(cudaError_t code, const char *file, int line, bool abort=
 }
 #endif
 
+#if  TENSORFLOW_USE_ROCM
+#include<hip/hip_runtime.h>
+
+#define hipErrcheck(res) { hipAssert((res), __FILE__, __LINE__); }
+inline void hipAssert(hipError_t code, const char *file, int line, bool abort=true)
+{
+    if (code != hipSuccess)
+    {
+        fprintf(stderr,"hip assert: %s %s %d\n", hipGetErrorString(code), file, line);
+        if (abort) exit(code);
+    }
+}
+#endif //TENSORFLOW_USE_ROCM
 
 static 
 std::vector<int> cum_sum (const std::vector<int32> & n_sel) {
@@ -217,6 +230,20 @@ init (const std::string & model, const int & gpu_rank, const std::string & file_
     graph::SetDefaultDevice(str, &graph_def);
   }
   #endif // GOOGLE_CUDA
+
+  #if TENSORFLOW_USE_ROCM
+  hipGetDeviceCount(&gpu_num); // check current device environment
+  if (gpu_num > 0) {
+    options.config.set_allow_soft_placement(true);
+    options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.9);
+    options.config.mutable_gpu_options()->set_allow_growth(true);
+    hipErrcheck(hipSetDevice(gpu_rank % gpu_num));
+    std::string str = "/gpu:";
+    str += std::to_string(gpu_rank % gpu_num);
+    graph::SetDefaultDevice(str, &graph_def);
+  }
+  #endif // TENSORFLOW_USE_ROCM
+
   check_status (NewSession(options, &session));
   check_status (session->Create(graph_def));
   rcut = get_scalar<VALUETYPE>("descrpt_attr/rcut");
@@ -524,6 +551,10 @@ init (const std::vector<std::string> & models, const int & gpu_rank, const std::
   cudaGetDeviceCount(&gpu_num);
   #endif // GOOGLE_CUDA
 
+  #if TENSORFLOW_USE_ROCM
+  hipGetDeviceCount(&gpu_num);
+  #endif //TENSORFLOW_USE_ROCM
+
   SessionOptions options;
   options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
   options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
@@ -542,6 +573,16 @@ init (const std::vector<std::string> & models, const int & gpu_rank, const std::
   }
   #endif // GOOGLE_CUDA
 
+
+  #if TENSORFLOW_USE_ROCM
+  if (gpu_num > 0) {
+      options.config.set_allow_soft_placement(true);
+      options.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.9);
+      options.config.mutable_gpu_options()->set_allow_growth(true);
+      hipErrcheck(hipSetDevice(gpu_rank % gpu_num));
+  }
+  #endif // TENSORFLOW_USE_ROCM
+
   for (unsigned ii = 0; ii < numb_models; ++ii) {
     if (gpu_num > 0) {
       std::string str = "/gpu:";

source/api_cc/tests/CMakeLists.txt

@@ -48,33 +48,36 @@ find_package(Threads)
 # find openmp
 find_package(OpenMP)
 if (OPENMP_FOUND)
-    set (CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
-    set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
+  set (CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
+  set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
 endif()
 
-# define USE_CUDA_TOOLKIT
-if (DEFINED USE_CUDA_TOOLKIT)
-  if (USE_CUDA_TOOLKIT)
-    find_package(CUDA REQUIRED)
-  else()
-    message(STATUS "Will not build nv GPU support")
-  endif()
-else()
-  find_package(CUDA QUIET)  
-  if (CUDA_FOUND)
-    set(USE_CUDA_TOOLKIT TRUE)
-    message(STATUS "Found CUDA in ${CUDA_TOOLKIT_ROOT_DIR}, build nv GPU support")
-  else()
-    set(USE_CUDA_TOOLKIT FALSE)
-    message(STATUS "No cuda support found, will not build nv GPU support")
-  endif()
+# Devices that have both ROCM and CUDA are not currently supported
+if (USE_ROCM_TOOLKIT AND USE_CUDA_TOOLKIT)
+  message (FATAL_ERROR "Devices that have both ROCM and CUDA are not currently supported")
 endif()
 
+# define USE_CUDA_TOOLKIT
 if (USE_CUDA_TOOLKIT)
-  add_definitions("-D GOOGLE_CUDA")
+  find_package(CUDA REQUIRED)
+  add_definitions("-DGOOGLE_CUDA")
+  message(STATUS "Found CUDA in ${CUDA_TOOLKIT_ROOT_DIR}, build nv GPU support")
   include_directories(${CUDA_INCLUDE_DIRS})
   add_subdirectory(${LIB_BASE_DIR}/src/cuda cuda_binary_dir)
-endif()
+else()
+  message(STATUS "Will not build nv GPU support")
+endif(USE_CUDA_TOOLKIT)
+
+#define USE_ROCM_TOOLKIT
+if (USE_ROCM_TOOLKIT)
+  find_package(ROCM REQUIRED)
+  add_definitions("-DTENSORFLOW_USE_ROCM")
+  add_compile_definitions(__HIP_PLATFORM_HCC__)
+  include_directories(${ROCM_INCLUDE_DIRS})
+  add_subdirectory(${LIB_BASE_DIR}/src/rocm rocm_binary_dir)
+else()
+  message(STATUS "Will not build AMD GPU support")
+endif (USE_ROCM_TOOLKIT)
 
 file(GLOB TEST_SRC test_*.cc)
 add_executable( runUnitTests ${TEST_SRC} )
@@ -93,9 +96,12 @@ endif()
 
 if (USE_CUDA_TOOLKIT)
   target_link_libraries(runUnitTests gtest gtest_main ${libname} ${apiname} ${opname} pthread ${TensorFlow_LIBRARY} deepmd_op_cuda coverage_config)
+elseif(USE_ROCM_TOOLKIT)
+  target_link_libraries(runUnitTests gtest gtest_main ${libname} ${apiname} ${opname} pthread ${TensorFlow_LIBRARY} deepmd_op_rocm coverage_config)
 else()
   target_link_libraries(runUnitTests gtest gtest_main ${libname} ${apiname} ${opname} pthread ${TensorFlow_LIBRARY} coverage_config)
 endif()
+
 add_test( runUnitTests runUnitTests )
 
 find_package(GTest)

source/cmake/FindROCM.cmake

@@ -0,0 +1,80 @@
+# Input:
+# ROCM_ROOT 
+#
+# Output:
+# ROCM_FOUND        
+# ROCM_INCLUDE_DIRS 
+# 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 ()
+
+# includes
+find_path (ROCM_INCLUDE_DIRS
+  NAMES 
+  hip/hip_runtime.h
+  rocprim/rocprim.hpp
+  hipcub/hipcub.hpp
+  PATHS ${ROCM_search_PATHS} 
+  PATH_SUFFIXES "include"
+  NO_DEFAULT_PATH
+  )
+if (NOT ROCM_INCLUDE_DIRS AND ROCM_FIND_REQUIRED)
+  message(FATAL_ERROR 
+    "Not found 'hip' or 'rocprim' or 'hipcub' directory in path '${ROCM_search_PATHS}' "
+    "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 
+  FindHIP.cmake
+  PATHS ${ROCM_search_PATHS} 
+  PATH_SUFFIXES "hip/cmake"
+  NO_DEFAULT_PATH
+  )
+
+if (NOT HIP_CMAKE AND ROCM_FIND_REQUIRED)
+  message(FATAL_ERROR 
+    "Not found 'FindHIP.cmake' file in path '${ROCM_search_PATHS}' "
+    "You can manually set the ROCM install path by -DROCM_ROOT ")
+endif ()
+
+list (APPEND CMAKE_MODULE_PATH ${HIP_CMAKE})
+find_package(HIP) 
+
+# define the output variable
+if (ROCM_INCLUDE_DIRS AND ROCM_LIBRARIES AND HIP_CMAKE)
+  set(ROCM_FOUND TRUE)
+else ()
+  set(ROCM_FOUND FALSE)
+endif ()
+
+# print message
+if (NOT ROCM_FIND_QUIETLY)
+  message(STATUS "Found ROCM: ${ROCM_INCLUDE_DIRS}, ${ROCM_LIBRARIES}, ${HIP_CMAKE}"
+    " in ${ROCM_search_PATHS}, build AMD GPU support")
+endif ()
+
+unset(ROCM_search_PATHS)

source/lib/CMakeLists.txt

@@ -1,23 +1,33 @@
 # libmd
-
 set (libname ${LIB_DEEPMD})
 
 if (USE_CUDA_TOOLKIT)
   include_directories("${CUDA_INCLUDE_DIRS}")
 endif()
 
+if (USE_ROCM_TOOLKIT)
+  include_directories("${ROCM_INCLUDE_DIRS}")
+endif()
+
 file(GLOB LIB_SRC src/*.cc src/*.cpp)
 file(GLOB INC_SRC include/*.h ${CMAKE_CURRENT_BINARY_DIR}/version.h)
 
 add_library(${libname} SHARED ${LIB_SRC})
 
 if (USE_CUDA_TOOLKIT)
-  add_definitions("-D GOOGLE_CUDA")
+  add_definitions("-DGOOGLE_CUDA")
   add_subdirectory(src/cuda)
   set (EXTRA_LIBS ${EXTRA_LIBS} deepmd_op_cuda)
   target_link_libraries (${libname} ${CUDA_LIBRARIES} ${EXTRA_LIBS})
 endif()
 
+if (USE_ROCM_TOOLKIT)
+  add_definitions("-DTENSORFLOW_USE_ROCM")
+  add_subdirectory(src/rocm)
+  set (EXTRA_LIBS ${EXTRA_LIBS} deepmd_op_rocm)
+  target_link_libraries (${libname} ${ROCM_LIBRARIES} ${EXTRA_LIBS})
+endif()
+
 if(BUILD_PY_IF)
   install(TARGETS ${libname} DESTINATION deepmd/op/)
 endif(BUILD_PY_IF)