cmake: use modern HIP language (#2857)

Use modern `HIP` language instead of `hip_add_library`.
Use modern `find_package` to find `hip` and `hipcub` instead of
`FindROCM.cmake` (which is removed in this PR).

---------

Signed-off-by: Jinzhe Zeng <[email protected]>

Author: njzjz

backend/read_env.py

@@ -57,8 +57,8 @@ def get_argument_from_env() -> Tuple[str, list, list, dict, str]:
         if rocm_root:
             cmake_args.append(f"-DCMAKE_HIP_COMPILER_ROCM_ROOT:STRING={rocm_root}")
         hipcc_flags = os.environ.get("HIP_HIPCC_FLAGS")
-        if hipcc_flags:
-            cmake_args.append(f"-DHIP_HIPCC_FLAGS:STRING={hipcc_flags}")
+        if hipcc_flags is not None:
+            os.environ["HIPFLAGS"] = os.environ.get("HIPFLAGS", "") + " " + hipcc_flags
     else:
         raise RuntimeError("Unsupported DP_VARIANT option: %s" % dp_variant)
 

doc/install/install-from-source.md

@@ -79,6 +79,7 @@ One may set the following environment variables before executing `pip`:
 | TENSORFLOW_ROOT       | Path                   | Detected automatically | The path to TensorFlow Python library. By default the installer only finds TensorFlow under user site-package directory (`site.getusersitepackages()`) or system site-package directory (`sysconfig.get_path("purelib")`) due to limitation of [PEP-517](https://peps.python.org/pep-0517/). If not found, the latest TensorFlow (or the environment variable `TENSORFLOW_VERSION` if given) from PyPI will be built against.|
 | DP_ENABLE_NATIVE_OPTIMIZATION | 0, 1           | 0             | Enable compilation optimization for the native machine's CPU type. Do not enable it if generated code will run on different CPUs. |
 | CMAKE_ARGS             | str                   | -             | Additional CMake arguments |
+| &lt;LANG&gt;FLAGS (`<LANG>`=`CXX`, `CUDA` or `HIP`)   | str            | -             | Default compilation flags to be used when compiling `<LANG>` files. See [CMake documentation](https://cmake.org/cmake/help/latest/variable/CMAKE_LANG_FLAGS.html). |
 
 To test the installation, one should first jump out of the source directory
 ```
@@ -193,7 +194,8 @@ One may add the following arguments to `cmake`:
 | -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.|
-| -DENABLE_NATIVE_OPTIMIZATION       | `TRUE` or `FALSE` | `FALSE`       | Enable compilation optimization for the native machine's CPU type. Do not enable it if generated code will run on different CPUs. |
+| -DENABLE_NATIVE_OPTIMIZATION=&lt;value&gt;       | `TRUE` or `FALSE` | `FALSE`       | Enable compilation optimization for the native machine's CPU type. Do not enable it if generated code will run on different CPUs. |
+| -DCMAKE_&lt;LANG&gt;_FLAGS=&lt;value&gt; (`<LANG>`=`CXX`, `CUDA` or `HIP`)   | str            | -             | Default compilation flags to be used when compiling `<LANG>` files. See [CMake documentation](https://cmake.org/cmake/help/latest/variable/CMAKE_LANG_FLAGS.html). |
 
 If the CMake has been executed successfully, then run the following make commands to build the package:
 ```bash

source/CMakeLists.txt

@@ -111,10 +111,14 @@ endif(USE_CUDA_TOOLKIT)
 # define USE_ROCM_TOOLKIT
 if(USE_ROCM_TOOLKIT)
   cmake_minimum_required(VERSION 3.21)
-  find_package(ROCM REQUIRED)
+  include(CMakeDetermineHIPCompiler)
+  list(APPEND CMAKE_PREFIX_PATH ${CMAKE_HIP_COMPILER_ROCM_ROOT})
+  find_package(hip REQUIRED)
+  find_package(hipcub REQUIRED)
   add_definitions("-DTENSORFLOW_USE_ROCM")
-  add_compile_definitions(__HIP_PLATFORM_HCC__)
-  message(STATUS "Found ROCM in ${ROCM_ROOT}, build AMD GPU support")
+  message(
+    STATUS
+      "Found ROCM in ${CMAKE_HIP_COMPILER_ROCM_ROOT}, build AMD GPU support")
   set(DP_VARIANT "rocm")
 else()
   message(STATUS "Will not build AMD GPU support")

source/cmake/FindROCM.cmake

@@ -24,10 +24,12 @@ if(USE_ROCM_TOOLKIT)
   add_definitions("-DTENSORFLOW_USE_ROCM")
   add_subdirectory(src/gpu)
   set(EXTRA_LIBS ${EXTRA_LIBS} deepmd_op_rocm)
-  target_link_libraries(${libname} INTERFACE ${ROCM_LIBRARIES} ${EXTRA_LIBS})
+  # to define __HIP_PLATFORM_AMD__ in hip_runtime.h
+  target_link_libraries(${libname} PUBLIC hip::host)
+  target_link_libraries(${libname} INTERFACE ${EXTRA_LIBS})
   # gpu_rocm.h
   target_include_directories(
-    ${libname} PUBLIC $<BUILD_INTERFACE:${ROCM_INCLUDE_DIRS}>
+    ${libname} PUBLIC $<BUILD_INTERFACE:${HIP_INCLUDE_DIRS}>
                       $<INSTALL_INTERFACE:include>)
 endif()
 

source/lib/CMakeLists.txt

@@ -50,26 +50,29 @@ elseif(USE_ROCM_TOOLKIT)
   cmake_minimum_required(VERSION 3.21)
   # project name
   project(deepmd_op_rocm)
+  enable_language(HIP)
   set(GPU_LIB_NAME deepmd_op_rocm)
   set(CMAKE_LINK_WHAT_YOU_USE TRUE)
 
   # set c++ version c++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})
+  message(STATUS "HIP major version is " ${hip_VERSION_MAJOR})
 
-  set(HIP_HIPCC_FLAGS -fno-gpu-rdc; -fPIC --std=c++14 ${HIP_HIPCC_FLAGS}
-  )# --amdgpu-target=gfx906
-  if(HIP_VERSION VERSION_LESS 3.5.1)
-    set(HIP_HIPCC_FLAGS -hc; ${HIP_HIPCC_FLAGS})
+  set(CMAKE_HIP_FLAGS -fno-gpu-rdc ${CMAKE_HIP_FLAGS}) # --amdgpu-target=gfx906
+  if(hip_VERSION VERSION_LESS 3.5.1)
+    set(CMAKE_HIP_FLAGS -hc ${CMAKE_HIP_FLAGS})
   endif()
 
   file(GLOB SOURCE_FILES "*.cu")
 
-  hip_add_library(${GPU_LIB_NAME} SHARED ${SOURCE_FILES})
+  add_library(${GPU_LIB_NAME} SHARED ${SOURCE_FILES})
+  set_source_files_properties(${SOURCE_FILES} PROPERTIES LANGUAGE HIP)
+  # -fpic
+  set_property(TARGET ${GPU_LIB_NAME} PROPERTY POSITION_INDEPENDENT_CODE ON)
+  target_link_libraries(${GPU_LIB_NAME} PRIVATE hip::hipcub)
 
 endif()