Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Tpl: adding rocSOLVER for Lapack #2034

Merged
merged 2 commits into from
Nov 15, 2023
Merged

Tpl: adding rocSOLVER for Lapack #2034

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
24c73c8
LAPACK: adding rocsolver TPL
lucbv Nov 9, 2023
c06b8db
Lapack: change according to Brian's review
lucbv Nov 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion cmake/Dependencies.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm, I wonder if this is breaking with Trilinos/TriBITS in #2036 ?

TEST_OPTIONAL_TPLS yaml-cpp
)
# NOTE: If you update names in LIB_OPTIONAL_TPLS above, make sure to map those names in
Expand Down
2 changes: 2 additions & 0 deletions cmake/KokkosKernels_config.h.in
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
9 changes: 9 additions & 0 deletions cmake/Modules/FindTPLROCSOLVER.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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()
12 changes: 12 additions & 0 deletions cmake/kokkoskernels_features.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
8 changes: 6 additions & 2 deletions cmake/kokkoskernels_tpls.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 ## _")
Expand Down Expand Up @@ -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)
Expand Down
7 changes: 7 additions & 0 deletions common/src/KokkosKernels_PrintConfiguration.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 0 additions & 7 deletions lapack/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 #
Expand Down
44 changes: 27 additions & 17 deletions lapack/impl/KokkosLapack_gesv_spec.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 };
};
Expand All @@ -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 }; \
};

Expand All @@ -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,
lucbv marked this conversation as resolved.
Show resolved Hide resolved
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) "
Expand All @@ -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>
Expand Down
44 changes: 28 additions & 16 deletions lapack/src/KokkosLapack_gesv.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
lucbv marked this conversation as resolved.
Show resolved Hide resolved
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,
Expand Down Expand Up @@ -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);
}
}

Expand Down
Loading