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

Jump to bottom

Kokkos Kernels: Add MAGMA TPL support for GESV on HIP #13436

Open
wants to merge 6 commits into
base: develop
Choose a base branch
from
+135 −59
Open

Kokkos Kernels: Add MAGMA TPL support for GESV on HIP #13436

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f82d513
Changes for supporting MAGMA GESV on HIP
vqd8a Sep 9, 2024
201dfb8
Merge branch 'develop' into kk-magmma-gesv-hip
vqd8a Sep 12, 2024
7b3be99
Merge branch 'develop' into kk-magmma-gesv-hip
vqd8a Sep 24, 2024
a94e8b8
Modifications for check validity of output views sizes in svd
vqd8a Oct 1, 2024
e18732a
Merge branch 'develop' into kk-magmma-gesv-hip
vqd8a Oct 1, 2024
e4120fc
Merge branch 'develop' into kk-magmma-gesv-hip
vqd8a Oct 3, 2024
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
16 changes: 16 additions & 0 deletions packages/kokkos-kernels/lapack/src/KokkosLapack_gesv.hpp
View file
Open in desktop
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
18 changes: 16 additions & 2 deletions packages/kokkos-kernels/lapack/src/KokkosLapack_svd.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,20 +127,34 @@ void svd(const ExecutionSpace& space, const char jobu[], const char jobvt[], con
is_extent_invalid = true;
os << "KokkosLapack::svd: S has extent " << S.extent(0) << ", instead of " << rankA << ".\n";
}
if ((jobu[0] == 'A') || (jobu[0] == 'a') || (jobu[0] == 'S') || (jobu[0] == 's')) {
if ((jobu[0] == 'A') || (jobu[0] == 'a')) {
if (U.extent_int(0) != m || U.extent_int(1) != m) {
is_extent_invalid = true;
os << "KokkosLapack::svd: U has extents (" << U.extent(0) << ", " << U.extent(1) << ") instead of (" << m << ", "
<< m << ").\n";
}
}
if ((jobvt[0] == 'A') || (jobvt[0] == 'a') || (jobvt[0] == 'S') || (jobvt[0] == 's')) {
if ((jobu[0] == 'S') || (jobu[0] == 's')) {
if (U.extent_int(0) != m || U.extent_int(1) != std::min(m,n)) {
is_extent_invalid = true;
os << "KokkosLapack::svd: U has extents (" << U.extent(0) << ", " << U.extent(1) << ") instead of (" << m << ", "
<< std::min(m,n) << ").\n";
}
}
if ((jobvt[0] == 'A') || (jobvt[0] == 'a')) {
if (Vt.extent_int(0) != n || Vt.extent_int(1) != n) {
is_extent_invalid = true;
os << "KokkosLapack::svd: V has extents (" << Vt.extent(0) << ", " << Vt.extent(1) << ") instead of (" << n
<< ", " << n << ").\n";
}
}
if ((jobvt[0] == 'S') || (jobvt[0] == 's')) {
if (Vt.extent_int(0) != std::min(m,n) || Vt.extent_int(1) != n) {
is_extent_invalid = true;
os << "KokkosLapack::svd: V has extents (" << Vt.extent(0) << ", " << Vt.extent(1) << ") instead of (" << std::min(m,n)
<< ", " << n << ").\n";
}
}
if (is_extent_invalid) {
KokkosKernels::Impl::throw_runtime_exception(os.str());
}
Expand Down
37 changes: 21 additions & 16 deletions packages/kokkos-kernels/lapack/tpls/KokkosLapack_gesv_tpl_spec_avail.hpp
View file
Open in desktop
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 packages/kokkos-kernels/lapack/tpls/KokkosLapack_gesv_tpl_spec_decl.hpp
View file
Open in desktop
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
Loading