GPU CMake: Multiple updates

Author: davidrohr

dependencies/FindO2GPU.cmake

@@ -10,9 +10,9 @@
 # or submit itself to any jurisdiction.
 
 # NOTE!!!! - Whenever this file is changed, move it over to alidist/resources
-# FindO2GPU.cmake Version 18
+# FindO2GPU.cmake Version 19
 
-set(CUDA_COMPUTETARGET_DEFAULT_FULL 80-real 86-real 89-real 120-real 75-virtual)
+set(CUDA_COMPUTETARGET_DEFAULT_FULL 80-real;86-real;89-real;120-real;75-virtual)
 set(HIP_AMDGPUTARGET_DEFAULT_FULL gfx906;gfx908)
 set(CUDA_COMPUTETARGET_DEFAULT_MINIMAL 75-virtual)
 set(HIP_AMDGPUTARGET_DEFAULT_MINIMAL gfx906)
@@ -179,6 +179,11 @@ if(ENABLE_CUDA)
   endif()
   set(CMAKE_CUDA_STANDARD ${CMAKE_CXX_STANDARD})
   set(CMAKE_CUDA_STANDARD_REQUIRED TRUE)
+  if (DEFINED ENV{O2_GPU_CUDA_HOME})
+    set(CMAKE_CUDA_COMPILER "$ENV{O2_GPU_CUDA_HOME}/bin/nvcc")
+  elseif (DEFINED ENV{CUDA_PATH})
+    set(CMAKE_CUDA_COMPILER "$ENV{CUDA_PATH}/bin/nvcc")
+  endif()
   include(CheckLanguage)
   check_language(CUDA)
   if (NOT ENABLE_CUDA STREQUAL "AUTO")
@@ -310,18 +315,30 @@ if(ENABLE_HIP)
     set(CMAKE_HIP_ARCHITECTURES "${HIP_AMDGPUTARGET}")
     set(GPU_TARGETS "${HIP_AMDGPUTARGET}")
   endif()
-  if(NOT "$ENV{CMAKE_PREFIX_PATH}" MATCHES "rocm" AND NOT CMAKE_PREFIX_PATH MATCHES "rocm" AND EXISTS "/opt/rocm/lib/cmake/")
+  if (DEFINED ENV{O2_GPU_ROCM_HOME})
+    list(PREPEND CMAKE_PREFIX_PATH "$ENV{O2_GPU_ROCM_HOME}/lib/cmake")
+  elseif (DEFINED ENV{ROCM_PATH})
+    list(PREPEND CMAKE_PREFIX_PATH "$ENV{ROCM_PATH}/lib/cmake")
+  elseif(NOT "$ENV{CMAKE_PREFIX_PATH}" MATCHES "rocm|ROCm" AND NOT CMAKE_PREFIX_PATH MATCHES "rocm|ROCm" AND EXISTS "/opt/rocm/lib/cmake/")
     list(APPEND CMAKE_PREFIX_PATH "/opt/rocm/lib/cmake")
   endif()
   # TODO: Use ROCm folder as provided by alidist-gpu recipe
-  if("$ENV{CMAKE_PREFIX_PATH}" MATCHES "rocm" OR CMAKE_PREFIX_PATH MATCHES "rocm")
+  if("$ENV{CMAKE_PREFIX_PATH}" MATCHES "rocm|ROCm" OR CMAKE_PREFIX_PATH MATCHES "rocm|ROCm")
+    if (DEFINED ENV{O2_GPU_ROCM_HOME})
+      set(CMAKE_HIP_COMPILER "$ENV{O2_GPU_ROCM_HOME}/llvm/bin/clang++")
+      set(TMP_ROCM_DIR "$ENV{O2_GPU_ROCM_HOME}")
+    elseif (DEFINED ENV{ROCM_PATH})
+      set(CMAKE_HIP_COMPILER "$ENV{ROCM_PATH}/llvm/bin/clang++")
+      set(TMP_ROCM_DIR "$ENV{ROCM_PATH}")
+    else()
+      set(TMP_ROCM_DIR_LIST "${CMAKE_PREFIX_PATH}:$ENV{CMAKE_PREFIX_PATH}")
+      string(REPLACE ":" ";" TMP_ROCM_DIR_LIST "${TMP_ROCM_DIR_LIST}")
+      list(FILTER TMP_ROCM_DIR_LIST INCLUDE REGEX /rocm/lib/cmake|/ROCm/lib/cmake)
+      list(POP_FRONT TMP_ROCM_DIR_LIST TMP_ROCM_DIR)
+      get_filename_component(TMP_ROCM_DIR ${TMP_ROCM_DIR}/../../ ABSOLUTE)
+    endif()
     set(CMAKE_HIP_STANDARD ${CMAKE_CXX_STANDARD})
     set(CMAKE_HIP_STANDARD_REQUIRED TRUE)
-    set(TMP_ROCM_DIR_LIST "${CMAKE_PREFIX_PATH}:$ENV{CMAKE_PREFIX_PATH}")
-    string(REPLACE ":" ";" TMP_ROCM_DIR_LIST "${TMP_ROCM_DIR_LIST}")
-    list(FILTER TMP_ROCM_DIR_LIST INCLUDE REGEX /rocm/lib/cmake)
-    list(POP_FRONT TMP_ROCM_DIR_LIST TMP_ROCM_DIR)
-    get_filename_component(TMP_ROCM_DIR ${TMP_ROCM_DIR}/../../ ABSOLUTE)
     if (NOT DEFINED CMAKE_HIP_COMPILER)
       set(CMAKE_HIP_COMPILER "${TMP_ROCM_DIR}/llvm/bin/clang++")
       if(NOT EXISTS ${CMAKE_HIP_COMPILER})
@@ -362,7 +379,7 @@ if(ENABLE_HIP)
       enable_language(HIP)
     endif()
   elseif(NOT ENABLE_HIP STREQUAL "AUTO")
-    message(FATAL_ERROR "HIP requested, but CMAKE_PREFIX_PATH env variable does not contain rocm folder!")
+    message(FATAL_ERROR "HIP requested, but CMAKE_PREFIX_PATH (${CMAKE_PREFIX_PATH} env variable does not contain rocm folder!")
   endif()
   if(hip_FOUND AND NOT hip_VERSION VERSION_GREATER_EQUAL "6.3")
     set(hip_FOUND 0)