Merge pull request kokkos#2034 from lucbv/tpl_rocsolver

Tpl: adding rocSOLVER for Lapack

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,48 @@ 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);
+#if defined(KOKKOSKERNELS_ENABLE_TPL_MAGMA)
+  if constexpr (!std::is_same_v<ExecutionSpace, Kokkos::Cuda>) {
+    static_assert(
+        Kokkos::SpaceAccessibility<ExecutionSpace,
+                                   typename IPIVV::memory_space>::accessible);
+  }
+#else
+  static_assert(
+      Kokkos::SpaceAccessibility<ExecutionSpace,
+                                 typename IPIVV::memory_space>::accessible);
+#endif
   static_assert(Kokkos::is_view<AMatrix>::value,
                 "KokkosLapack::gesv: A must be a Kokkos::View.");
   static_assert(Kokkos::is_view<BXMV>::value,
@@ -137,12 +157,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);
   }
 }