Implement rocSOLVER dsyevd for MAGMA build

CMakeLists.txt

@@ -253,6 +253,33 @@ if (CUDA_FOUND)
    include_directories(${CUDA_INCLUDE_DIRS})
 endif()
 
+if(BML_MAGMA)
+  find_package(hip)
+endif()
+
+if(BML_ROCSOLVER)
+  find_package(rocblas REQUIRED)
+  find_package(rocsolver REQUIRED)
+endif()
+
+if (hip_FOUND)
+   message("HIP libraries: ${hip_LIBRARIES}")
+   list(APPEND LINK_LIBRARIES ${hip_LIBRARIES})
+   include_directories(${hip_INCLUDE_DIRS})
+endif()
+
+if (rocblas_FOUND)
+   message("ROCBLAS libraries: ${rocblas_LIBRARIES}")
+   list(APPEND LINK_LIBRARIES ${rocblas_LIBRARIES})
+   include_directories(${rocblas_INCLUDE_DIRS})
+endif()
+
+if (rocsolver_FOUND)
+   message("ROCSOLVER libraries: ${rocsolver_LIBRARIES}")
+   list(APPEND LINK_LIBRARIES ${rocsolver_LIBRARIES})
+   include_directories(${rocsolver_INCLUDE_DIRS})
+endif()
+
 set(BML_OPENCL FALSE CACHE BOOL "Whether to compiler with OpenCL support")
 if(BML_OPENCL)
   include(FindOpenCL)
@@ -261,9 +288,11 @@ endif()
 SET(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
 
 set(BML_MAGMA FALSE CACHE BOOL "Whether to use MAGMA library")
+
 if(BML_MAGMA)
   FIND_PACKAGE(MAGMA)
 endif()
+
 if (MAGMA_FOUND)
   add_definitions(-DBML_USE_MAGMA)
   message(STATUS
@@ -284,6 +313,10 @@ if (MAGMA_FOUND)
     message(STATUS "Will use cuSOLVER")
     add_definitions(-DBML_USE_CUSOLVER)
   endif()
+  if(BML_ROCSOLVER)
+    message(STATUS "Will use rocSOLVER")
+    add_definitions(-DBML_USE_ROCSOLVER)
+  endif()
 endif()
 
 if (CUBLAS_FOUND)
@@ -344,6 +377,9 @@ if(MAGMA_FOUND)
   if(BML_CUSOLVER)
     message(STATUS "cuSOLVER: ${CUDA_cusolver_LIBRARY}")
   endif()
+  if(BML_ROCSOLVER)
+      message(STATUS "rocSOLVER: ${ROCM_rocsolver_LIBRARIES}")
+  endif()
 endif()
 
 #check existence of required math and linear algebra functions

build.sh

@@ -83,6 +83,7 @@ EOF
     echo "BML_CUDA               Build with CUDA             (default is ${BML_CUDA})"
     echo "BML_MAGMA              Build with MAGMA            (default is ${BML_MAGMA})"
     echo "BML_CUSOLVER           Build with cuSOLVER         (default is ${BML_CUSOLVER})"
+    echo "BML_ROCSOLVER          Build with rocSOLVER        (default is ${BML_ROCSOLVER})"
     echo "BML_XSMM               Build with XSMM             (default is ${BML_XSMM})"
     echo "BML_SCALAPACK          Build with SCALAPACK        (default is ${BML_SCALAPACK})"
     echo "SCALAPACK_LIBRARIES    ScaLapack libraries         (default is ${SCALAPACK_LIBRARIES})"
@@ -129,6 +130,7 @@ set_defaults() {
     : ${BML_CUDA:=no}
     : ${BML_MAGMA:=no}
     : ${BML_CUSOLVER:=no}
+    : ${BML_ROCSOLVER:=no}
     : ${BML_XSMM:=no}
     : ${BML_SCALAPACK:=no}
     : ${BML_ELLBLOCK_MEMPOOL:=no}
@@ -215,6 +217,7 @@ configure() {
         -DBML_CUDA="${BML_CUDA}" \
         -DBML_MAGMA="${BML_MAGMA}" \
         -DBML_CUSOLVER="${BML_CUSOLVER}" \
+        -DBML_ROCSOLVER="${BML_ROCSOLVER}" \
         -DBML_XSMM="${BML_XSMM}" \
         -DBML_SCALAPACK="${BML_SCALAPACK}" \
         -DBML_ELLBLOCK_MEMPOOL="${BML_ELLBLOCK_MEMPOOL}" \

scripts/build_crusher_magma_gcc.sh

@@ -15,12 +15,19 @@ export BML_OPENMP=${BML_OPENMP:=yes}
 export INSTALL_DIR=${INSTALL_DIR:="${MY_PATH}/install"}
 export BML_MAGMA=${BML_MAGMA:=yes}
 export MAGMA_ROOT=${OLCF_MAGMA_ROOT}
+export BML_ROCSOLVER=${BML_ROCSOLVER:=no}
 export BML_TESTING=${BML_TESTING:=yes}
 export CMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE:=Release}
 export CMAKE_Fortran_FLAGS=${CMAKE_Fortran_FLAGS:="-g -ffree-form -ffree-line-length-200"}
 export CMAKE_C_FLAGS=${CMAKE_C_FLAGS:="-g -I${HIP_PATH}/include"}
 export EXTRA_LINK_FLAGS=${EXTRA_LINK_FLAGS:=""}
 
+if [[ -z "$CMAKE_PREFIX_PATH" ]]; then
+    export CMAKE_PREFIX_PATH=${ROCM_PATH}
+else
+    export CMAKE_PREFIX_PATH="${ROCM_PATH};${CMAKE_PREFIX_PATH}"
+fi
+
 ./build.sh configure
 
 pushd build

scripts/build_spock_magma_gcc.sh

@@ -15,12 +15,19 @@ export BML_OPENMP=${BML_OPENMP:=yes}
 export INSTALL_DIR=${INSTALL_DIR:="${MY_PATH}/install"}
 export BML_MAGMA=${BML_MAGMA:=yes}
 export MAGMA_ROOT=${OLCF_MAGMA_ROOT}
+export BML_ROCSOLVER=${BML_ROCSOLVER:=no}
 export BML_TESTING=${BML_TESTING:=yes}
 export CMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE:=Release}
 export CMAKE_Fortran_FLAGS=${CMAKE_Fortran_FLAGS:="-g -ffree-form -ffree-line-length-200"}
 export CMAKE_C_FLAGS=${CMAKE_C_FLAGS:="-g -I${HIP_PATH}/include"}
 export EXTRA_LINK_FLAGS=${EXTRA_LINK_FLAGS:=""}
 
+if [[ -z "$CMAKE_PREFIX_PATH" ]]; then
+    export CMAKE_PREFIX_PATH=${ROCM_PATH}
+else
+    export CMAKE_PREFIX_PATH="${ROCM_PATH};${CMAKE_PREFIX_PATH}"
+fi
+
 ./build.sh configure
 
 pushd build

scripts/setenv_crusher_magma_gcc.sh

@@ -4,3 +4,5 @@ module load PrgEnv-gnu
 module load magma
 module load cmake
 module load openblas
+module load rocm
+

scripts/setenv_spock_magma_gcc.sh

@@ -4,3 +4,4 @@ module load PrgEnv-gnu
 module load magma
 module load cmake
 module load openblas
+module load rocm

src/C-interface/dense/bml_diagonalize_dense.c

@@ -14,6 +14,11 @@
 #include <cuda_runtime.h>
 #include <cusolverDn.h>
 #endif
+#ifdef BML_USE_ROCSOLVER
+#include <hip/hip_runtime_api.h>
+#include <rocblas.h>
+#include <rocsolver.h>
+#endif
 #else
 #include "../lapack.h"
 #endif
@@ -189,6 +194,53 @@ bml_diagonalize_dense_double_real(
 
     if (cusolverH)
         cusolverDnDestroy(cusolverH);
+#else
+#ifdef BML_USE_ROCSOLVER
+    // See https://rocsolver.readthedocs.io/_/downloads/en/latest/pdf/
+    // create cusolver/cublas handle
+    rocblas_handle rocblasH = NULL;
+    rocblas_status rocblasS = rocblas_create_handle(&rocblasH);
+    assert(rocblas_status_success == rocblasS);
+
+    // allocate memory for eigenvalues
+    double *d_W = NULL;
+    hipError_t hipStat = hipMalloc((void **) &d_W, sizeof(double) * A->N);
+    assert(hipSuccess == hipStat);
+
+    // compute eigenvalues and eigenvectors
+    rocblas_evect evect = rocblas_evect_original;
+    rocblas_fill uplo = rocblas_fill_lower;
+
+    // allocate working space of syevd
+    double *d_work = NULL;
+    hipStat = hipMalloc((void **) &d_work, sizeof(double) * A->N * A->N);
+    assert(hipSuccess == hipStat);
+
+    // solve
+    rocblas_int *devInfo = NULL;
+    hipStat = hipMalloc((void **) &devInfo, sizeof(rocblas_int));
+    assert(hipSuccess == hipStat);
+
+    rocblasS =
+        rocsolver_dsyevd(rocblasH, evect, uplo, A->N, evecs, A->ld, d_W,
+                         d_work, devInfo);
+    hipStat = hipDeviceSynchronize();
+    assert(rocblas_status_success == rocblasS);
+    assert(hipSuccess == hipStat);
+
+    // copy eigenvalues to CPU
+    hipStat =
+        hipMemcpy(typed_eigenvalues, d_W, sizeof(double) * A->N,
+                  hipMemcpyDeviceToHost);
+    assert(hipSuccess == hipStat);
+
+    // free resources
+    hipFree(d_W);
+    hipFree(devInfo);
+    hipFree(d_work);
+
+    if (rocblasH)
+        rocblas_destroy_handle(rocblasH);
 #else // MAGMA
     int nb = magma_get_ssytrd_nb(A->N);
 
@@ -222,7 +274,7 @@ bml_diagonalize_dense_double_real(
     //    printf("norm = %le\n", norm);
     //}
 #endif
-
+#endif
     // transpose eigenvactors matrix on GPU
     A_matrix = (double *) eigenvectors->matrix;
     magmablas_dtranspose(A->N, A->N, evecs, A->ld,