Add support for ROCm 3.1.0

ROCm shared objects were re-organized into subfolders in version
3.1.0, one folder per library (hip, hcc, rocfft, rocprim, etc.).
The CMake logic now supports ROCm 3.0.x and 3.1.0.

CMakeLists.txt

@@ -187,13 +187,20 @@ if(WITH_CUDA)
       message(STATUS "Found HIP compiler: ${HIP_HIPCC_EXECUTABLE}")
       set(CUDA 1)
       set(HIP 1)
-      list(APPEND HIP_HCC_FLAGS "-I${HIP_ROOT_DIR}/include -Wno-c99-designator -Wno-macro-redefined -Wno-duplicate-decl-specifier -std=c++14")
+      list(APPEND HIP_HCC_FLAGS "-I/opt/rocm/include -I${HIP_ROOT_DIR}/include -Wno-c99-designator -Wno-macro-redefined -Wno-duplicate-decl-specifier -std=c++14")
       list(APPEND HIP_HCC_FLAGS "-pedantic -Wall -Wextra -Wno-sign-compare -Wno-unused-function -Wno-unused-variable -Wno-unused-parameter -Wno-missing-braces -Wno-gnu-anonymous-struct -Wno-nested-anon-types -Wno-gnu-zero-variadic-macro-arguments")
       if(WARNINGS_ARE_ERRORS)
         list(APPEND HIP_HCC_FLAGS "-Werror")
       endif()
 
-      find_library(ROCFFT_LIB name "rocfft" PATHS "${HIP_ROOT_DIR}/lib")
+      set(ROCFFT_PATH "${HIP_ROOT_DIR}")
+      if(EXISTS "/opt/rocm/.info/version-dev")
+        file(READ "/opt/rocm/.info/version-dev" ROCM_VERSION)
+        if(ROCM_VERSION VERSION_GREATER_EQUAL 3.1.0)
+          set(ROCFFT_PATH "/opt/rocm/rocfft")
+        endif()
+      endif()
+      find_library(ROCFFT_LIB name "rocfft" PATHS "${ROCFFT_PATH}/lib")
 
       function(add_gpu_library)
         hip_add_library(${ARGV})