MegEngine/cmake/rocm.cmake

if(NOT DEFINED HIP_PATH)
  if(NOT DEFINED ENV{HIP_PATH})
    set(HIP_PATH
        "/opt/rocm/hip"
        CACHE PATH "Path to which HIP has been installed")
  else()
    set(HIP_PATH
        $ENV{HIP_PATH}
        CACHE PATH "Path to which HIP has been installed")
  endif()
endif()
set(CMAKE_MODULE_PATH "${HIP_PATH}/cmake" ${CMAKE_MODULE_PATH})
find_package(HIP QUIET)
if(HIP_FOUND)
  message(STATUS "Found HIP: " ${HIP_VERSION})
else()
  message(
    FATAL_ERROR
      "Could not find HIP. Ensure that HIP is either installed in /opt/rocm/hip or the variable HIP_PATH is set to point to the right location."
  )
endif()

if(${HIP_VERSION} VERSION_LESS 3.0)
  message(FATAL_ERROR "ROCM version needed 3. Please update ROCM.")
endif()

macro(hipconfig_get_option variable option)
  if(NOT DEFINED ${variable})
    execute_process(COMMAND ${HIP_HIPCONFIG_EXECUTABLE} ${option}
                    OUTPUT_VARIABLE ${variable})
  endif()
endmacro()

hipconfig_get_option(HIP_COMPILER "--compiler")
hipconfig_get_option(HIP_CPP_CONFIG "--cpp_config")

separate_arguments(HIP_CPP_CONFIG)

foreach(hip_config_item ${HIP_CPP_CONFIG})
  foreach(macro_name "__HIP_PLATFORM_HCC__" "__HIP_ROCclr__")
    if(${hip_config_item} STREQUAL "-D${macro_name}=")
      set(HIP_CPP_DEFINE "${HIP_CPP_DEFINE}#define ${macro_name}\n")
      set(HIP_CPP_UNDEFINE
          "${HIP_CPP_UNDEFINE}\
                    #ifdef ${macro_name}\n#undef ${macro_name}\n\
                    #else\n#error\n\
                    #endif\n")
    elseif(${hip_config_item} STREQUAL "-D${macro_name}")
      set(HIP_CPP_DEFINE "${HIP_CPP_DEFINE}#define ${macro_name} 1\n")
      set(HIP_CPP_UNDEFINE
          "${HIP_CPP_UNDEFINE}\
                    #ifdef ${macro_name}\n#undef ${macro_name}\n\
                    #else\n#error\n\
                    #endif\n")
    endif()
  endforeach()
endforeach()

message(STATUS "Using HIP compiler ${HIP_COMPILER}")

if(${HIP_COMPILER} STREQUAL "hcc")
  set(MGE_ROCM_LIBS hip_hcc)
  message(
    WARNING "hcc is not well supported, please modify link.txt to link with hipcc")
elseif(${HIP_COMPILER} STREQUAL "clang")
  set(MGE_ROCM_LIBS amdhip64)
endif()

list(APPEND MGE_ROCM_LIBS amdocl64 MIOpen rocblas rocrand)

set(HIP_INCLUDE_DIR ${HIP_ROOT_DIR}/../include)
set(HIP_LIBRARY_DIR ${HIP_ROOT_DIR}/../lib)

function(find_rocm_library name dirname include library)
  find_path(
    ${name}_LIBRARY_DIR
    NAMES ${library}
    HINTS "${${name}_ROOT_DIR}" "${HIP_ROOT_DIR}/../${dirname}"
    PATH_SUFFIXES lib lib/x86_64
    DOC "Path to ${name} library directory")

  if(${${name}_LIBRARY_DIR} MATCHES "NOTFOUND$")
    message(FATAL_ERROR "Can not find ${name} library")
  endif()

  find_path(
    ${name}_INCLUDE_DIR
    NAMES ${include}
    HINTS "${${name}_ROOT_DIR}" "${HIP_ROOT_DIR}/../${dirname}"
    PATH_SUFFIXES include
    DOC "Path to ${name} include directory")

  if(${name}_INCLUDE_DIR MATCHES "NOTFOUND$")
    message(FATAL_ERROR "Can not find ${name} include")
  endif()
  message(DEBUG "Found lib ${${name}_LIBRARY_DIR}, include ${${name}_INCLUDE_DIR}")
endfunction()

find_rocm_library(MIOPEN miopen miopen libMIOpen.so)
find_rocm_library(ROCBLAS rocblas rocblas.h librocblas.so)
find_rocm_library(ROCRAND rocrand rocrand.h librocrand.so)
find_rocm_library(AMDOCL opencl CL libamdocl64.so)