LAPACK: adding rocsolver TPL

Adding the necessary CMake logic and TPL layer to support rocsolver
for LAPACK. Enabling the TPL in gesv and updating gesv test to run
by default the more common configurations and only run specific ones
when the associated TPL (MAGMA) is enabled.

CMakeLists.txt

@@ -377,6 +377,7 @@ ELSE()
   KOKKOSKERNELS_LINK_TPL(kokkoskernels PUBLIC CUSPARSE)
   KOKKOSKERNELS_LINK_TPL(kokkoskernels PUBLIC ROCBLAS)
   KOKKOSKERNELS_LINK_TPL(kokkoskernels PUBLIC ROCSPARSE)
+  KOKKOSKERNELS_LINK_TPL(kokkoskernels PUBLIC ROCSOLVER)
   KOKKOSKERNELS_LINK_TPL(kokkoskernels PUBLIC METIS)
   KOKKOSKERNELS_LINK_TPL(kokkoskernels PUBLIC ARMPL)
   KOKKOSKERNELS_LINK_TPL(kokkoskernels PUBLIC MAGMA)

cmake/Dependencies.cmake

@@ -1,6 +1,6 @@
 TRIBITS_PACKAGE_DEFINE_DEPENDENCIES(
         LIB_REQUIRED_PACKAGES Kokkos
-        LIB_OPTIONAL_TPLS quadmath MKL BLAS LAPACK CUSPARSE METIS SuperLU Cholmod CUBLAS ROCBLAS ROCSPARSE
+        LIB_OPTIONAL_TPLS quadmath MKL BLAS LAPACK METIS SuperLU Cholmod CUBLAS CUSPARSE ROCBLAS ROCSPARSE ROCSOLVER
         TEST_OPTIONAL_TPLS yaml-cpp
 )
 # NOTE: If you update names in LIB_OPTIONAL_TPLS above, make sure to map those names in

cmake/KokkosKernels_config.h.in

@@ -138,6 +138,8 @@
 #cmakedefine KOKKOSKERNELS_ENABLE_TPL_ROCBLAS
 /* ROCSPARSE */
 #cmakedefine KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE
+/* ROCSOLVER */
+#cmakedefine KOKKOSKERNELS_ENABLE_TPL_ROCSOLVER
 
 #cmakedefine KOKKOSKERNELS_ENABLE_SUPERNODAL_SPTRSV
 

cmake/Modules/FindTPLROCSOLVER.cmake

@@ -0,0 +1,9 @@
+# LBV: 11/08/2023: This file follows the partern of FindTPLROCBLAS.cmake/FindTPLROCSPARSE.cmake
+FIND_PACKAGE(ROCSOLVER)
+if(TARGET roc::rocsolver)
+  SET(TPL_ROCSOLVER_IMPORTED_NAME roc::rocsolver)
+  SET(TPL_IMPORTED_NAME roc::rocsolver)
+  ADD_LIBRARY(KokkosKernels::ROCSOLVER ALIAS roc::rocsolver)
+ELSE()
+  MESSAGE(FATAL_ERROR "Package ROCSOLVER requested but not found")
+ENDIF()

cmake/kokkoskernels_features.cmake

@@ -27,3 +27,15 @@ IF (KOKKOSKERNELS_ENABLE_TPL_BLAS OR KOKKOSKERNELS_ENABLE_TPL_MKL OR KOKKOSKERNE
   INCLUDE(CheckHostBlasReturnComplex.cmake)
   CHECK_HOST_BLAS_RETURN_COMPLEX(KOKKOSKERNELS_TPL_BLAS_RETURN_COMPLEX)
 ENDIF()
+
+# ==================================================================
+# Lapack requirements
+# ==================================================================
+
+IF (KOKKOSKERNELS_ENABLE_TPL_ROCSOLVER AND NOT KOKKOSKERNELS_ENABLE_TPL_ROCBLAS AND NOT KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE)
+  MESSAGE(FATAL_ERROR "rocSOLVER requires rocBLAS and rocSPARSE, please reconfigure with KOKKOSKERNELS_ENABLE_TPL_ROCBLAS:BOOL=ON and KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE:BOOL=ON.")
+ELSEIF (KOKKOSKERNELS_ENABLE_TPL_ROCSOLVER AND NOT KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE)
+  MESSAGE(FATAL_ERROR "rocSOLVER requires rocSPARSE, please reconfigure with KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE:BOOL=ON.")
+ELSEIF (KOKKOSKERNELS_ENABLE_TPL_ROCSOLVER AND NOT KOKKOSKERNELS_ENABLE_TPL_ROCBLAS)
+  MESSAGE(FATAL_ERROR "rocSOLVER requires rocBLAS, please reconfigure with KOKKOSKERNELS_ENABLE_TPL_ROCBLAS:BOOL=ON.")
+ENDIF()

cmake/kokkoskernels_tpls.cmake

@@ -460,15 +460,18 @@ KOKKOSKERNELS_ADD_OPTION(NO_DEFAULT_ROCM_TPLS OFF BOOL "Whether ROCM TPLs should
 # Unlike CUDA, ROCm does not automatically install these TPLs
 SET(ROCBLAS_DEFAULT   OFF)
 SET(ROCSPARSE_DEFAULT OFF)
+SET(ROCSOLVER_DEFAULT OFF)
 # Since the default is OFF we do not really need this piece of logic here.
 # IF(KOKKOSKERNELS_NO_DEFAULT_ROCM_TPLS)
 #   SET(ROCBLAS_DEFAULT   OFF)
 #   SET(ROCSPARSE_DEFAULT OFF)
 # ENDIF()
 KOKKOSKERNELS_ADD_TPL_OPTION(ROCBLAS   ${ROCBLAS_DEFAULT}     "Whether to enable ROCBLAS"
-  DEFAULT_DOCSTRING "ON if HIP-enabled Kokkos, otherwise OFF")
+  DEFAULT_DOCSTRING "OFF even if HIP-enabled Kokkos")
 KOKKOSKERNELS_ADD_TPL_OPTION(ROCSPARSE ${ROCSPARSE_DEFAULT}   "Whether to enable ROCSPARSE"
-  DEFAULT_DOCSTRING "ON if HIP-enabled Kokkos, otherwise OFF")
+  DEFAULT_DOCSTRING "OFF even if HIP-enabled Kokkos")
+KOKKOSKERNELS_ADD_TPL_OPTION(ROCSOLVER ${ROCSOLVER_DEFAULT}   "Whether to enable ROCSOLVER"
+  DEFAULT_DOCSTRING "OFF even if HIP-enabled Kokkos")
 
 IF (KOKKOSKERNELS_ENABLE_TPL_MAGMA)
     IF (F77_BLAS_MANGLE STREQUAL "(name,NAME) name ## _")
@@ -507,6 +510,7 @@ IF (NOT KOKKOSKERNELS_HAS_TRILINOS)
   KOKKOSKERNELS_IMPORT_TPL(MAGMA)
   KOKKOSKERNELS_IMPORT_TPL(ROCBLAS)
   KOKKOSKERNELS_IMPORT_TPL(ROCSPARSE)
+  KOKKOSKERNELS_IMPORT_TPL(ROCSOLVER)
 ELSE ()
   IF (Trilinos_ENABLE_SuperLU5_API)
     SET(HAVE_KOKKOSKERNELS_SUPERLU5_API TRUE)

common/src/KokkosKernels_PrintConfiguration.hpp

@@ -110,6 +110,13 @@ inline void print_enabled_tpls(std::ostream& os) {
   os << "  "
      << "KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE: no\n";
 #endif
+#ifdef KOKKOSKERNELS_ENABLE_TPL_ROCSOLVER
+  os << "  "
+     << "KOKKOSKERNELS_ENABLE_TPL_ROCSOLVER: yes\n";
+#else
+  os << "  "
+     << "KOKKOSKERNELS_ENABLE_TPL_ROCOLVER: no\n";
+#endif
 #ifdef KOKKOSKERNELS_ENABLE_TPL_METIS
   os << "KOKKOSKERNELS_ENABLE_TPL_METIS: yes\n";
 #else

lapack/CMakeLists.txt

@@ -34,13 +34,6 @@ IF (KOKKOSKERNELS_ENABLE_TPL_CUSOLVER)
   )
 ENDIF()
 
-# Include rocm lapack TPL source file
-IF (KOKKOSKERNELS_ENABLE_TPL_ROCSOLVER)
-  LIST(APPEND SOURCES
-   lapack/tpls/KokkosLapack_Rocm_tpl.cpp
-  )
-ENDIF()
-
 ##################
 #                #
 # ETI generation #

lapack/impl/KokkosLapack_gesv_spec.hpp

@@ -28,7 +28,7 @@
 namespace KokkosLapack {
 namespace Impl {
 // Specialization struct which defines whether a specialization exists
-template <class AVT, class BVT>
+template <class ExecutionSpace, class AVT, class BVT, class IPIVV>
 struct gesv_eti_spec_avail {
   enum : bool { value = false };
 };
@@ -46,12 +46,16 @@ struct gesv_eti_spec_avail {
                                          EXEC_SPACE_TYPE, MEM_SPACE_TYPE) \
   template <>                                                             \
   struct gesv_eti_spec_avail<                                             \
+      EXEC_SPACE_TYPE,                                                    \
       Kokkos::View<SCALAR_TYPE **, LAYOUT_TYPE,                           \
                    Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,       \
-                   Kokkos::MemoryTraits<Kokkos::Unmanaged> >,             \
+                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>,              \
       Kokkos::View<SCALAR_TYPE **, LAYOUT_TYPE,                           \
                    Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,       \
-                   Kokkos::MemoryTraits<Kokkos::Unmanaged> > > {          \
+                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>,              \
+      Kokkos::View<int *, LAYOUT_TYPE,                                    \
+                   Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,       \
+                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>> {            \
     enum : bool { value = true };                                         \
   };
 
@@ -65,20 +69,24 @@ namespace Impl {
 // Unification layer
 /// \brief Implementation of KokkosLapack::gesv.
 
-template <class AMatrix, class BXMV, class IPIVV,
-          bool tpl_spec_avail = gesv_tpl_spec_avail<AMatrix, BXMV>::value,
-          bool eti_spec_avail = gesv_eti_spec_avail<AMatrix, BXMV>::value>
+template <class ExecutionSpace, class AMatrix, class BXMV, class IPIVV,
+          bool tpl_spec_avail =
+              gesv_tpl_spec_avail<ExecutionSpace, AMatrix, BXMV, IPIVV>::value,
+          bool eti_spec_avail =
+              gesv_eti_spec_avail<ExecutionSpace, AMatrix, BXMV, IPIVV>::value>
 struct GESV {
-  static void gesv(const AMatrix &A, const BXMV &B, const IPIVV &IPIV);
+  static void gesv(const ExecutionSpace &space, const AMatrix &A, const BXMV &B,
+                   const IPIVV &IPIV);
 };
 
 #if !defined(KOKKOSKERNELS_ETI_ONLY) || KOKKOSKERNELS_IMPL_COMPILE_LIBRARY
 //! Full specialization of gesv for multi vectors.
 // Unification layer
-template <class AMatrix, class BXMV, class IPIVV>
-struct GESV<AMatrix, BXMV, IPIVV, false, KOKKOSKERNELS_IMPL_COMPILE_LIBRARY> {
-  static void gesv(const AMatrix & /* A */, const BXMV & /* B */,
-                   const IPIVV & /* IPIV */) {
+template <class ExecutionSpace, class AMatrix, class BXMV, class IPIVV>
+struct GESV<ExecutionSpace, AMatrix, BXMV, IPIVV, false,
+            KOKKOSKERNELS_IMPL_COMPILE_LIBRARY> {
+  static void gesv(const ExecutionSpace & /* space */, const AMatrix & /* A */,
+                   const BXMV & /* B */, const IPIVV & /* IPIV */) {
     // NOTE: Might add the implementation of KokkosLapack::gesv later
     throw std::runtime_error(
         "No fallback implementation of GESV (general LU factorization & solve) "
@@ -100,31 +108,33 @@ struct GESV<AMatrix, BXMV, IPIVV, false, KOKKOSKERNELS_IMPL_COMPILE_LIBRARY> {
 #define KOKKOSLAPACK_GESV_ETI_SPEC_DECL(SCALAR_TYPE, LAYOUT_TYPE,        \
                                         EXEC_SPACE_TYPE, MEM_SPACE_TYPE) \
   extern template struct GESV<                                           \
+      EXEC_SPACE_TYPE,                                                   \
       Kokkos::View<SCALAR_TYPE **, LAYOUT_TYPE,                          \
                    Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,      \
-                   Kokkos::MemoryTraits<Kokkos::Unmanaged> >,            \
+                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>,             \
       Kokkos::View<SCALAR_TYPE **, LAYOUT_TYPE,                          \
                    Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,      \
-                   Kokkos::MemoryTraits<Kokkos::Unmanaged> >,            \
+                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>,             \
       Kokkos::View<int *, LAYOUT_TYPE,                                   \
                    Kokkos::Device<Kokkos::DefaultHostExecutionSpace,     \
                                   Kokkos::HostSpace>,                    \
-                   Kokkos::MemoryTraits<Kokkos::Unmanaged> >,            \
+                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>,             \
       false, true>;
 
 #define KOKKOSLAPACK_GESV_ETI_SPEC_INST(SCALAR_TYPE, LAYOUT_TYPE,        \
                                         EXEC_SPACE_TYPE, MEM_SPACE_TYPE) \
   template struct GESV<                                                  \
+      EXEC_SPACE_TYPE,                                                   \
       Kokkos::View<SCALAR_TYPE **, LAYOUT_TYPE,                          \
                    Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,      \
-                   Kokkos::MemoryTraits<Kokkos::Unmanaged> >,            \
+                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>,             \
       Kokkos::View<SCALAR_TYPE **, LAYOUT_TYPE,                          \
                    Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,      \
-                   Kokkos::MemoryTraits<Kokkos::Unmanaged> >,            \
+                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>,             \
       Kokkos::View<int *, LAYOUT_TYPE,                                   \
                    Kokkos::Device<Kokkos::DefaultHostExecutionSpace,     \
                                   Kokkos::HostSpace>,                    \
-                   Kokkos::MemoryTraits<Kokkos::Unmanaged> >,            \
+                   Kokkos::MemoryTraits<Kokkos::Unmanaged>>,             \
       false, true>;
 
 #include <KokkosLapack_gesv_tpl_spec_decl.hpp>

lapack/src/KokkosLapack_gesv.hpp

@@ -34,28 +34,40 @@ namespace KokkosLapack {
 
 /// \brief Solve the dense linear equation system A*X = B.
 ///
+/// \tparam ExecutionSpace the space where the kernel will run.
 /// \tparam AMatrix Input matrix/Output LU, as a 2-D Kokkos::View.
 /// \tparam BXMV Input (right-hand side)/Output (solution) (multi)vector, as a
-/// 1-D or 2-D Kokkos::View. \tparam IPIVV Output pivot indices, as a 1-D
-/// Kokkos::View
+/// 1-D or 2-D Kokkos::View.
+/// \tparam IPIVV Output pivot indices, as a 1-D Kokkos::View
 ///
 /// \param A [in,out] On entry, the N-by-N matrix to be solved. On exit, the
 /// factors L and U from
 ///   the factorization A = P*L*U; the unit diagonal elements of L are not
 ///   stored.
 /// \param B [in,out] On entry, the right hand side (multi)vector B. On exit,
-/// the solution (multi)vector X. \param IPIV [out] On exit, the pivot indices
-/// (for partial pivoting). If the View extents are zero and
-///   its data pointer is NULL, pivoting is not used.
+/// the solution (multi)vector X.
+/// \param IPIV [out] On exit, the pivot indices (for partial pivoting).
+/// If the View extents are zero and its data pointer is NULL, pivoting is not
+/// used.
 ///
-template <class AMatrix, class BXMV, class IPIVV>
-void gesv(const AMatrix& A, const BXMV& B, const IPIVV& IPIV) {
-  // NOTE: Currently, KokkosLapack::gesv only supports for MAGMA TPL and LAPACK
-  // TPL.
-  //       MAGMA TPL should be enabled to call the MAGMA GPU interface for
-  //       device views LAPACK TPL should be enabled to call the LAPACK
-  //       interface for host views
+template <class ExecutionSpace, class AMatrix, class BXMV, class IPIVV>
+void gesv(const ExecutionSpace& space, const AMatrix& A, const BXMV& B,
+          const IPIVV& IPIV) {
+  // NOTE: Currently, KokkosLapack::gesv only supports LAPACK, MAGMA and
+  // rocSOLVER TPLs.
+  //       MAGMA/rocSOLVER TPL should be enabled to call the MAGMA/rocSOLVER GPU
+  //       interface for device views LAPACK TPL should be enabled to call the
+  //       LAPACK interface for host views
 
+  static_assert(
+      Kokkos::SpaceAccessibility<ExecutionSpace,
+                                 typename AMatrix::memory_space>::accessible);
+  static_assert(
+      Kokkos::SpaceAccessibility<ExecutionSpace,
+                                 typename BXMV::memory_space>::accessible);
+  static_assert(
+      Kokkos::SpaceAccessibility<ExecutionSpace,
+                                 typename IPIVV::memory_space>::accessible);
   static_assert(Kokkos::is_view<AMatrix>::value,
                 "KokkosLapack::gesv: A must be a Kokkos::View.");
   static_assert(Kokkos::is_view<BXMV>::value,
@@ -137,12 +149,12 @@ void gesv(const AMatrix& A, const BXMV& B, const IPIVV& IPIV) {
 
   if (BXMV::rank == 1) {
     auto B_i = BXMV_Internal(B.data(), B.extent(0), 1);
-    KokkosLapack::Impl::GESV<AMatrix_Internal, BXMV_Internal,
-                             IPIVV_Internal>::gesv(A_i, B_i, IPIV_i);
+    KokkosLapack::Impl::GESV<ExecutionSpace, AMatrix_Internal, BXMV_Internal,
+                             IPIVV_Internal>::gesv(space, A_i, B_i, IPIV_i);
   } else {  // BXMV::rank == 2
     auto B_i = BXMV_Internal(B.data(), B.extent(0), B.extent(1));
-    KokkosLapack::Impl::GESV<AMatrix_Internal, BXMV_Internal,
-                             IPIVV_Internal>::gesv(A_i, B_i, IPIV_i);
+    KokkosLapack::Impl::GESV<ExecutionSpace, AMatrix_Internal, BXMV_Internal,
+                             IPIVV_Internal>::gesv(space, A_i, B_i, IPIV_i);
   }
 }