Skip to content

Commit

Permalink
Add MAGMA TPL support for GESV on HIP backend (#2326)
Browse files Browse the repository at this point in the history 
* Small changes for MAGMA GESV on HIP

* Apply clang-format

* Relax eps to 1e-8 for multi-rhs tests

---------

Co-authored-by: Vinh Quang Dang (-EXP) <[email protected]>
  • Loading branch information
@vqd8a
vqd8a and Vinh Quang Dang (-EXP) authored Sep 6, 2024
1 parent f26fbca commit 2c4dd7e
Show file tree
Hide file tree
Showing 4 changed files with 119 additions and 57 deletions.
16 changes: 16 additions & 0 deletions lapack/src/KokkosLapack_gesv.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,9 +63,15 @@ void gesv(const ExecutionSpace& space, const AMatrix& A, const BXMV& B, const IP
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 defined(KOKKOS_ENABLE_CUDA)
if constexpr (!std::is_same_v<ExecutionSpace, Kokkos::Cuda>) {
static_assert(Kokkos::SpaceAccessibility<ExecutionSpace, typename IPIVV::memory_space>::accessible);
}
#elif defined(KOKKOS_ENABLE_HIP)
if constexpr (!std::is_same_v<ExecutionSpace, Kokkos::HIP>) {
static_assert(Kokkos::SpaceAccessibility<ExecutionSpace, typename IPIVV::memory_space>::accessible);
}
#endif
#else
static_assert(Kokkos::SpaceAccessibility<ExecutionSpace, typename IPIVV::memory_space>::accessible);
#endif
Expand Down Expand Up @@ -96,13 +102,23 @@ void gesv(const ExecutionSpace& space, const AMatrix& A, const BXMV& B, const IP
// Check for no pivoting case. Only MAGMA supports no pivoting interface
#ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA // have MAGMA TPL
#ifdef KOKKOSKERNELS_ENABLE_TPL_LAPACK // and have LAPACK TPL
#if defined(KOKKOS_ENABLE_CUDA)
if ((!std::is_same<typename AMatrix::device_type::memory_space, Kokkos::CudaSpace>::value) && (IPIV0 == 0) &&
(IPIV.data() == nullptr)) {
std::ostringstream os;
os << "KokkosLapack::gesv: IPIV: " << IPIV0 << ". "
<< "LAPACK TPL does not support no pivoting.";
KokkosKernels::Impl::throw_runtime_exception(os.str());
}
#elif defined(KOKKOS_ENABLE_HIP)
if ((!std::is_same<typename AMatrix::device_type::memory_space, Kokkos::HIPSpace>::value) && (IPIV0 == 0) &&
(IPIV.data() == nullptr)) {
std::ostringstream os;
os << "KokkosLapack::gesv: IPIV: " << IPIV0 << ". "
<< "LAPACK TPL does not support no pivoting.";
KokkosKernels::Impl::throw_runtime_exception(os.str());
}
#endif
#endif
#else // not have MAGMA TPL
#ifdef KOKKOSKERNELS_ENABLE_TPL_LAPACK // but have LAPACK TPL
Expand Down
37 changes: 21 additions & 16 deletions lapack/tpls/KokkosLapack_gesv_tpl_spec_avail.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,23 +52,28 @@ KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_LAPACK(Kokkos::complex<float>, Kokkos::LayoutLe

namespace KokkosLapack {
namespace Impl {
#define KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(SCALAR, LAYOUT, MEMSPACE) \
template <> \
struct gesv_tpl_spec_avail< \
Kokkos::Cuda, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEMSPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged> >, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEMSPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged> >, \
Kokkos::View<magma_int_t*, LAYOUT, Kokkos::Device<Kokkos::DefaultHostExecutionSpace, Kokkos::HostSpace>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged> > > { \
enum : bool { value = true }; \
#define KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(SCALAR, LAYOUT, EXECSPACE, MEMSPACE) \
template <> \
struct gesv_tpl_spec_avail< \
EXECSPACE, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXECSPACE, MEMSPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged> >, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXECSPACE, MEMSPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged> >, \
Kokkos::View<magma_int_t*, LAYOUT, Kokkos::Device<Kokkos::DefaultHostExecutionSpace, Kokkos::HostSpace>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged> > > { \
enum : bool { value = true }; \
};

KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(double, Kokkos::LayoutLeft, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(float, Kokkos::LayoutLeft, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::CudaSpace)
#if defined(KOKKOS_ENABLE_CUDA)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(double, Kokkos::LayoutLeft, Kokkos::Cuda, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(float, Kokkos::LayoutLeft, Kokkos::Cuda, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::Cuda, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::Cuda, Kokkos::CudaSpace)
#endif
#if defined(KOKKOS_ENABLE_HIP)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(double, Kokkos::LayoutLeft, Kokkos::HIP, Kokkos::HIPSpace)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(float, Kokkos::LayoutLeft, Kokkos::HIP, Kokkos::HIPSpace)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::HIP, Kokkos::HIPSpace)
KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::HIP, Kokkos::HIPSpace)
#endif
} // namespace Impl
} // namespace KokkosLapack
#endif // KOKKOSKERNELS_ENABLE_TPL_MAGMA
Expand Down
74 changes: 40 additions & 34 deletions lapack/tpls/KokkosLapack_gesv_tpl_spec_decl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -197,42 +197,48 @@ void magmaGesvWrapper(const ExecSpace& space, const AViewType& A, const BViewTyp
Kokkos::Profiling::popRegion();
}

#define KOKKOSLAPACK_GESV_MAGMA(SCALAR, LAYOUT, MEM_SPACE) \
template <> \
struct GESV<Kokkos::Cuda, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
Kokkos::View<magma_int_t*, LAYOUT, Kokkos::Device<Kokkos::DefaultHostExecutionSpace, Kokkos::HostSpace>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
true, \
gesv_eti_spec_avail<Kokkos::Cuda, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
Kokkos::View<magma_int_t*, LAYOUT, \
Kokkos::Device<Kokkos::DefaultHostExecutionSpace, Kokkos::HostSpace>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>>::value> { \
using AViewType = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>; \
using BViewType = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>; \
using PViewType = \
Kokkos::View<magma_int_t*, LAYOUT, Kokkos::Device<Kokkos::DefaultHostExecutionSpace, Kokkos::HostSpace>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>; \
\
static void gesv(const Kokkos::Cuda& space, const AViewType& A, const BViewType& B, const PViewType& IPIV) { \
magmaGesvWrapper(space, A, B, IPIV); \
} \
#define KOKKOSLAPACK_GESV_MAGMA(SCALAR, LAYOUT, EXEC_SPACE, MEM_SPACE) \
template <> \
struct GESV< \
EXEC_SPACE, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, MEM_SPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, MEM_SPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
Kokkos::View<magma_int_t*, LAYOUT, Kokkos::Device<Kokkos::DefaultHostExecutionSpace, Kokkos::HostSpace>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
true, \
gesv_eti_spec_avail< \
EXEC_SPACE, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
Kokkos::View<magma_int_t*, LAYOUT, Kokkos::Device<Kokkos::DefaultHostExecutionSpace, Kokkos::HostSpace>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>>::value> { \
using AViewType = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>; \
using BViewType = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, MEM_SPACE>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>; \
using PViewType = \
Kokkos::View<magma_int_t*, LAYOUT, Kokkos::Device<Kokkos::DefaultHostExecutionSpace, Kokkos::HostSpace>, \
Kokkos::MemoryTraits<Kokkos::Unmanaged>>; \
\
static void gesv(const EXEC_SPACE& space, const AViewType& A, const BViewType& B, const PViewType& IPIV) { \
magmaGesvWrapper(space, A, B, IPIV); \
} \
};

KOKKOSLAPACK_GESV_MAGMA(float, Kokkos::LayoutLeft, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_MAGMA(double, Kokkos::LayoutLeft, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_MAGMA(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_MAGMA(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::CudaSpace)

#if defined(KOKKOS_ENABLE_CUDA)
KOKKOSLAPACK_GESV_MAGMA(float, Kokkos::LayoutLeft, Kokkos::Cuda, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_MAGMA(double, Kokkos::LayoutLeft, Kokkos::Cuda, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_MAGMA(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::Cuda, Kokkos::CudaSpace)
KOKKOSLAPACK_GESV_MAGMA(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::Cuda, Kokkos::CudaSpace)
#endif
#if defined(KOKKOS_ENABLE_HIP)
KOKKOSLAPACK_GESV_MAGMA(float, Kokkos::LayoutLeft, Kokkos::HIP, Kokkos::HIPSpace)
KOKKOSLAPACK_GESV_MAGMA(double, Kokkos::LayoutLeft, Kokkos::HIP, Kokkos::HIPSpace)
KOKKOSLAPACK_GESV_MAGMA(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::HIP, Kokkos::HIPSpace)
KOKKOSLAPACK_GESV_MAGMA(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::HIP, Kokkos::HIPSpace)
#endif
} // namespace Impl
} // namespace KokkosLapack
#endif // KOKKOSKERNELS_ENABLE_TPL_MAGMA
Expand Down
Loading

0 comments on commit 2c4dd7e

Please sign in to comment.