// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project

#ifndef KOKKOSLAPACK_SVD_TPL_SPEC_DECL_HPP_
#define KOKKOSLAPACK_SVD_TPL_SPEC_DECL_HPP_

#include "KokkosKernels_Error.hpp"
#include "KokkosKernels_ArithTraits.hpp"

#include <iostream>

namespace KokkosLapack {
namespace Impl {
template <class ExecutionSpace, class AMatrix, class SVector, class UMatrix, class VMatrix>
inline void svd_print_specialization() {
#ifdef KOKKOSKERNELS_ENABLE_CHECK_SPECIALIZATION
#ifdef KOKKOSKERNELS_ENABLE_TPL_CUSOLVER
  if constexpr (std::is_same_v<ExecutionSpace, Kokkos::Cuda>) {
    printf(
        "KokkosLapack::svd<> TPL Cusolver specialization for < %s , %s, %s, %s "
        ">\n",
        typeid(AMatrix).name(), typeid(SVector).name(), typeid(UMatrix).name(), typeid(VMatrix).name());
  }
#endif
#endif
}
}  // namespace Impl
}  // namespace KokkosLapack

// LAPACK
#if defined(KOKKOSKERNELS_ENABLE_TPL_LAPACK) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL)
#include "KokkosLapack_Host_tpl.hpp"

namespace KokkosLapack {
namespace Impl {

template <class ExecutionSpace, class AMatrix, class SVector, class UMatrix, class VMatrix>
void lapackSvdWrapper(const ExecutionSpace& /* space */, const char jobu[], const char jobvt[], const AMatrix& A,
                      const SVector& S, const UMatrix& U, const VMatrix& Vt) {
  using memory_space = typename AMatrix::memory_space;
  using Scalar       = typename AMatrix::non_const_value_type;
  using Magnitude    = typename SVector::non_const_value_type;
  using ALayout_t    = typename AMatrix::array_layout;
  using ULayout_t    = typename UMatrix::array_layout;
  using VLayout_t    = typename VMatrix::array_layout;

  static_assert(std::is_same_v<ALayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: A needs to have a Kokkos::LayoutLeft");
  static_assert(std::is_same_v<ULayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: U needs to have a Kokkos::LayoutLeft");
  static_assert(std::is_same_v<VLayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: Vt needs to have a Kokkos::LayoutLeft");

  const int m    = A.extent_int(0);
  const int n    = A.extent_int(1);
  const int lda  = A.stride(1);
  const int ldu  = U.stride(1);
  const int ldvt = Vt.stride(1);

  int lwork = -1, info = 0;
  Kokkos::View<Magnitude*, memory_space> rwork("svd rwork buffer", 5 * Kokkos::min(m, n));
  Kokkos::View<Scalar*, memory_space> work("svd work buffer", 1);
  if constexpr (KokkosKernels::ArithTraits<Scalar>::is_complex) {
    HostLapack<std::complex<Magnitude>>::gesvd(
        jobu[0], jobvt[0], m, n, reinterpret_cast<std::complex<Magnitude>*>(A.data()), lda, S.data(),
        reinterpret_cast<std::complex<Magnitude>*>(U.data()), ldu,
        reinterpret_cast<std::complex<Magnitude>*>(Vt.data()), ldvt,
        reinterpret_cast<std::complex<Magnitude>*>(work.data()), lwork, rwork.data(), info);

    lwork = static_cast<int>(work(0).real());

    work = Kokkos::View<Scalar*, memory_space>("svd work buffer", lwork);
    HostLapack<std::complex<Magnitude>>::gesvd(
        jobu[0], jobvt[0], m, n, reinterpret_cast<std::complex<Magnitude>*>(A.data()), lda, S.data(),
        reinterpret_cast<std::complex<Magnitude>*>(U.data()), ldu,
        reinterpret_cast<std::complex<Magnitude>*>(Vt.data()), ldvt,
        reinterpret_cast<std::complex<Magnitude>*>(work.data()), lwork, rwork.data(), info);
  } else {
    HostLapack<Scalar>::gesvd(jobu[0], jobvt[0], m, n, A.data(), lda, S.data(), U.data(), ldu, Vt.data(), ldvt,
                              work.data(), lwork, rwork.data(), info);

    lwork = static_cast<int>(work(0));

    work = Kokkos::View<Scalar*, memory_space>("svd work buffer", lwork);
    HostLapack<Scalar>::gesvd(jobu[0], jobvt[0], m, n, A.data(), lda, S.data(), U.data(), ldu, Vt.data(), ldvt,
                              work.data(), lwork, rwork.data(), info);
  }
}

#define KOKKOSLAPACK_SVD_LAPACK(SCALAR, LAYOUT, EXEC_SPACE)                                                            \
  template <>                                                                                                          \
  struct SVD<EXEC_SPACE,                                                                                               \
             Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                             \
                          Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                    \
             Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                           \
                          Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>,     \
             Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                             \
                          Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                    \
             Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                             \
                          Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                    \
             true,                                                                                                     \
             svd_eti_spec_avail<                                                                                       \
                 EXEC_SPACE,                                                                                           \
                 Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                         \
                              Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                \
                 Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                       \
                              Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
                 Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                         \
                              Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                \
                 Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                         \
                              Kokkos::MemoryTraits<Kokkos::Unmanaged>>>::value> {                                      \
    using AMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                      \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                             \
    using SVector =                                                                                                    \
        Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                                \
                     Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>;          \
    using UMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                      \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                             \
    using VMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                      \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                             \
                                                                                                                       \
    static void svd(const EXEC_SPACE& space, const char jobu[], const char jobvt[], const AMatrix& A,                  \
                    const SVector& S, const UMatrix& U, const VMatrix& Vt) {                                           \
      Kokkos::Profiling::pushRegion("KokkosLapack::svd[TPL_LAPACK," #SCALAR "]");                                      \
      svd_print_specialization<EXEC_SPACE, AMatrix, SVector, UMatrix, VMatrix>();                                      \
                                                                                                                       \
      lapackSvdWrapper(space, jobu, jobvt, A, S, U, Vt);                                                               \
      Kokkos::Profiling::popRegion();                                                                                  \
    }                                                                                                                  \
  };

#if defined(KOKKOS_ENABLE_SERIAL)
KOKKOSLAPACK_SVD_LAPACK(float, Kokkos::LayoutLeft, Kokkos::Serial)
KOKKOSLAPACK_SVD_LAPACK(double, Kokkos::LayoutLeft, Kokkos::Serial)
KOKKOSLAPACK_SVD_LAPACK(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::Serial)
KOKKOSLAPACK_SVD_LAPACK(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::Serial)
#endif

#if defined(KOKKOS_ENABLE_OPENMP)
KOKKOSLAPACK_SVD_LAPACK(float, Kokkos::LayoutLeft, Kokkos::OpenMP)
KOKKOSLAPACK_SVD_LAPACK(double, Kokkos::LayoutLeft, Kokkos::OpenMP)
KOKKOSLAPACK_SVD_LAPACK(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::OpenMP)
KOKKOSLAPACK_SVD_LAPACK(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::OpenMP)
#endif

#if defined(KOKKOS_ENABLE_THREADS)
KOKKOSLAPACK_SVD_LAPACK(float, Kokkos::LayoutLeft, Kokkos::Threads)
KOKKOSLAPACK_SVD_LAPACK(double, Kokkos::LayoutLeft, Kokkos::Threads)
KOKKOSLAPACK_SVD_LAPACK(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::Threads)
KOKKOSLAPACK_SVD_LAPACK(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::Threads)
#endif

}  // namespace Impl
}  // namespace KokkosLapack
#endif  // KOKKOSKERNELS_ENABLE_TPL_LAPACK

#ifdef KOKKOSKERNELS_ENABLE_TPL_MKL
#include "mkl.h"

namespace KokkosLapack {
namespace Impl {

template <class ExecutionSpace, class AMatrix, class SVector, class UMatrix, class VMatrix>
void mklSvdWrapper(const ExecutionSpace& /* space */, const char jobu[], const char jobvt[], const AMatrix& A,
                   const SVector& S, const UMatrix& U, const VMatrix& Vt) {
  using memory_space = typename AMatrix::memory_space;
  using Scalar       = typename AMatrix::non_const_value_type;
  using Magnitude    = typename SVector::non_const_value_type;
  using ALayout_t    = typename AMatrix::array_layout;
  using ULayout_t    = typename UMatrix::array_layout;
  using VLayout_t    = typename VMatrix::array_layout;

  static_assert(std::is_same_v<ALayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: A needs to have a Kokkos::LayoutLeft");
  static_assert(std::is_same_v<ULayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: U needs to have a Kokkos::LayoutLeft");
  static_assert(std::is_same_v<VLayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: Vt needs to have a Kokkos::LayoutLeft");

  const lapack_int m    = A.extent_int(0);
  const lapack_int n    = A.extent_int(1);
  const lapack_int lda  = A.stride(1);
  const lapack_int ldu  = U.stride(1);
  const lapack_int ldvt = Vt.stride(1);

  Kokkos::View<Magnitude*, memory_space> rwork("svd rwork buffer", Kokkos::min(m, n) - 1);
  lapack_int ret = 0;
  if constexpr (std::is_same_v<Scalar, float>) {
    ret = LAPACKE_sgesvd(LAPACK_COL_MAJOR, jobu[0], jobvt[0], m, n, A.data(), lda, S.data(), U.data(), ldu, Vt.data(),
                         ldvt, rwork.data());
  }
  if constexpr (std::is_same_v<Scalar, double>) {
    ret = LAPACKE_dgesvd(LAPACK_COL_MAJOR, jobu[0], jobvt[0], m, n, A.data(), lda, S.data(), U.data(), ldu, Vt.data(),
                         ldvt, rwork.data());
  }
  if constexpr (std::is_same_v<Scalar, Kokkos::complex<float>>) {
    ret = LAPACKE_cgesvd(LAPACK_COL_MAJOR, jobu[0], jobvt[0], m, n, reinterpret_cast<lapack_complex_float*>(A.data()),
                         lda, S.data(), reinterpret_cast<lapack_complex_float*>(U.data()), ldu,
                         reinterpret_cast<lapack_complex_float*>(Vt.data()), ldvt, rwork.data());
  }
  if constexpr (std::is_same_v<Scalar, Kokkos::complex<double>>) {
    ret = LAPACKE_zgesvd(LAPACK_COL_MAJOR, jobu[0], jobvt[0], m, n, reinterpret_cast<lapack_complex_double*>(A.data()),
                         lda, S.data(), reinterpret_cast<lapack_complex_double*>(U.data()), ldu,
                         reinterpret_cast<lapack_complex_double*>(Vt.data()), ldvt, rwork.data());
  }

  if (ret != 0) {
    std::ostringstream os;
    os << "KokkosLapack::svd: MKL failed with return value: " << ret << "\n";
    KokkosKernels::Impl::throw_runtime_exception(os.str());
  }
}

#define KOKKOSLAPACK_SVD_MKL(SCALAR, LAYOUT, EXEC_SPACE)                                                               \
  template <>                                                                                                          \
  struct SVD<EXEC_SPACE,                                                                                               \
             Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                             \
                          Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                    \
             Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                           \
                          Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>,     \
             Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                             \
                          Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                    \
             Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                             \
                          Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                    \
             true,                                                                                                     \
             svd_eti_spec_avail<                                                                                       \
                 EXEC_SPACE,                                                                                           \
                 Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                         \
                              Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                \
                 Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                       \
                              Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
                 Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                         \
                              Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                \
                 Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                         \
                              Kokkos::MemoryTraits<Kokkos::Unmanaged>>>::value> {                                      \
    using AMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                      \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                             \
    using SVector =                                                                                                    \
        Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                                \
                     Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>;          \
    using UMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                      \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                             \
    using VMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<EXEC_SPACE, Kokkos::HostSpace>,                      \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                             \
                                                                                                                       \
    static void svd(const EXEC_SPACE& space, const char jobu[], const char jobvt[], const AMatrix& A,                  \
                    const SVector& S, const UMatrix& U, const VMatrix& Vt) {                                           \
      Kokkos::Profiling::pushRegion("KokkosLapack::svd[TPL_LAPACK," #SCALAR "]");                                      \
      svd_print_specialization<EXEC_SPACE, AMatrix, SVector, UMatrix, VMatrix>();                                      \
                                                                                                                       \
      mklSvdWrapper(space, jobu, jobvt, A, S, U, Vt);                                                                  \
      Kokkos::Profiling::popRegion();                                                                                  \
    }                                                                                                                  \
  };

#if defined(KOKKOS_ENABLE_SERIAL)
KOKKOSLAPACK_SVD_MKL(float, Kokkos::LayoutLeft, Kokkos::Serial)
KOKKOSLAPACK_SVD_MKL(double, Kokkos::LayoutLeft, Kokkos::Serial)
KOKKOSLAPACK_SVD_MKL(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::Serial)
KOKKOSLAPACK_SVD_MKL(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::Serial)
#endif

#if defined(KOKKOS_ENABLE_OPENMP)
KOKKOSLAPACK_SVD_MKL(float, Kokkos::LayoutLeft, Kokkos::OpenMP)
KOKKOSLAPACK_SVD_MKL(double, Kokkos::LayoutLeft, Kokkos::OpenMP)
KOKKOSLAPACK_SVD_MKL(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::OpenMP)
KOKKOSLAPACK_SVD_MKL(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::OpenMP)
#endif

#if defined(KOKKOS_ENABLE_THREADS)
KOKKOSLAPACK_SVD_MKL(float, Kokkos::LayoutLeft, Kokkos::Threads)
KOKKOSLAPACK_SVD_MKL(double, Kokkos::LayoutLeft, Kokkos::Threads)
KOKKOSLAPACK_SVD_MKL(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::Threads)
KOKKOSLAPACK_SVD_MKL(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::Threads)
#endif

}  // namespace Impl
}  // namespace KokkosLapack
#endif  // KOKKOSKERNELS_ENABLE_TPL_MKL

// CUSOLVER
#ifdef KOKKOSKERNELS_ENABLE_TPL_CUSOLVER
#include "KokkosLapack_cusolver.hpp"

namespace KokkosLapack {
namespace Impl {

template <class ExecutionSpace, class AMatrix, class SVector, class UMatrix, class VMatrix>
void cusolverSvdWrapper(const ExecutionSpace& space, const char jobu[], const char jobvt[], const AMatrix& A,
                        const SVector& S, const UMatrix& U, const VMatrix& Vt) {
  using memory_space = typename AMatrix::memory_space;
  using Scalar       = typename AMatrix::non_const_value_type;
  using Magnitude    = typename SVector::non_const_value_type;
  using ALayout_t    = typename AMatrix::array_layout;
  using ULayout_t    = typename UMatrix::array_layout;
  using VLayout_t    = typename VMatrix::array_layout;

  static_assert(std::is_same_v<ALayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: A needs to have a Kokkos::LayoutLeft");
  static_assert(std::is_same_v<ULayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: U needs to have a Kokkos::LayoutLeft");
  static_assert(std::is_same_v<VLayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: Vt needs to have a Kokkos::LayoutLeft");

  const int m    = A.extent_int(0);
  const int n    = A.extent_int(1);
  const int lda  = A.stride(1);
  const int ldu  = U.stride(1);
  const int ldvt = Vt.stride(1);

  int lwork = 0;
  Kokkos::View<int, memory_space> info("svd info");
  Kokkos::View<Magnitude*, memory_space> rwork("svd rwork buffer", Kokkos::min(m, n) - 1);

  CudaLapackSingleton& s = CudaLapackSingleton::singleton();
  KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(cusolverDnSetStream(s.handle, space.cuda_stream()));
  if constexpr (std::is_same_v<Scalar, float>) {
    KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(cusolverDnSgesvd_bufferSize(s.handle, m, n, &lwork));
    Kokkos::View<Scalar*, memory_space> work("svd work buffer", lwork);

    KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(cusolverDnSgesvd(s.handle, jobu[0], jobvt[0], m, n, A.data(), lda, S.data(),
                                                          U.data(), ldu, Vt.data(), ldvt, work.data(), lwork,
                                                          rwork.data(), info.data()));
  }
  if constexpr (std::is_same_v<Scalar, double>) {
    KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(cusolverDnDgesvd_bufferSize(s.handle, m, n, &lwork));
    Kokkos::View<Scalar*, memory_space> work("svd work buffer", lwork);

    KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(cusolverDnDgesvd(s.handle, jobu[0], jobvt[0], m, n, A.data(), lda, S.data(),
                                                          U.data(), ldu, Vt.data(), ldvt, work.data(), lwork,
                                                          rwork.data(), info.data()));
  }
  if constexpr (std::is_same_v<Scalar, Kokkos::complex<float>>) {
    KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(cusolverDnCgesvd_bufferSize(s.handle, m, n, &lwork));
    Kokkos::View<Scalar*, memory_space> work("svd work buffer", lwork);

    KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(
        cusolverDnCgesvd(s.handle, jobu[0], jobvt[0], m, n, reinterpret_cast<cuComplex*>(A.data()), lda, S.data(),
                         reinterpret_cast<cuComplex*>(U.data()), ldu, reinterpret_cast<cuComplex*>(Vt.data()), ldvt,
                         reinterpret_cast<cuComplex*>(work.data()), lwork, rwork.data(), info.data()));
  }
  if constexpr (std::is_same_v<Scalar, Kokkos::complex<double>>) {
    KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(cusolverDnZgesvd_bufferSize(s.handle, m, n, &lwork));
    Kokkos::View<Scalar*, memory_space> work("svd work buffer", lwork);

    KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(cusolverDnZgesvd(
        s.handle, jobu[0], jobvt[0], m, n, reinterpret_cast<cuDoubleComplex*>(A.data()), lda, S.data(),
        reinterpret_cast<cuDoubleComplex*>(U.data()), ldu, reinterpret_cast<cuDoubleComplex*>(Vt.data()), ldvt,
        reinterpret_cast<cuDoubleComplex*>(work.data()), lwork, rwork.data(), info.data()));
  }
  KOKKOSLAPACK_IMPL_CUSOLVER_SAFE_CALL(cusolverDnSetStream(s.handle, NULL));
}

#define KOKKOSLAPACK_SVD_CUSOLVER(SCALAR, LAYOUT, MEM_SPACE)                                                        \
  template <>                                                                                                       \
  struct SVD<Kokkos::Cuda,                                                                                          \
             Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>,                                \
                          Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                 \
             Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                        \
                          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<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>,                                \
                          Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                 \
             true,                                                                                                  \
             svd_eti_spec_avail<                                                                                    \
                 Kokkos::Cuda,                                                                                      \
                 Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>,                            \
                              Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                             \
                 Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                    \
                              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<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>,                            \
                              Kokkos::MemoryTraits<Kokkos::Unmanaged>>>::value> {                                   \
    using AMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>,                         \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                          \
    using SVector = Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                 \
                                 Kokkos::Device<Kokkos::Cuda, MEM_SPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>; \
    using UMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>,                         \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                          \
    using VMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEM_SPACE>,                         \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                          \
                                                                                                                    \
    static void svd(const Kokkos::Cuda& space, const char jobu[], const char jobvt[], const AMatrix& A,             \
                    const SVector& S, const UMatrix& U, const VMatrix& Vt) {                                        \
      Kokkos::Profiling::pushRegion("KokkosLapack::svd[TPL_CUSOLVER," #SCALAR "]");                                 \
      svd_print_specialization<Kokkos::Cuda, AMatrix, SVector, UMatrix, VMatrix>();                                 \
                                                                                                                    \
      cusolverSvdWrapper(space, jobu, jobvt, A, S, U, Vt);                                                          \
      Kokkos::Profiling::popRegion();                                                                               \
    }                                                                                                               \
  };

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

#if defined(KOKKOSKERNELS_INST_MEMSPACE_CUDAUVMSPACE)
KOKKOSLAPACK_SVD_CUSOLVER(float, Kokkos::LayoutLeft, Kokkos::CudaUVMSpace)
KOKKOSLAPACK_SVD_CUSOLVER(double, Kokkos::LayoutLeft, Kokkos::CudaUVMSpace)
KOKKOSLAPACK_SVD_CUSOLVER(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::CudaUVMSpace)
KOKKOSLAPACK_SVD_CUSOLVER(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::CudaUVMSpace)
#endif

}  // namespace Impl
}  // namespace KokkosLapack
#endif  // KOKKOSKERNELS_ENABLE_TPL_CUSOLVER

// ROCSOLVER
#ifdef KOKKOSKERNELS_ENABLE_TPL_ROCSOLVER
#include <KokkosBlas_tpl_spec.hpp>
#include <rocsolver/rocsolver.h>

namespace KokkosLapack {
namespace Impl {

template <class ExecutionSpace, class AMatrix, class SVector, class UMatrix, class VMatrix>
void rocsolverSvdWrapper(const ExecutionSpace& space, const char jobu[], const char jobvt[], const AMatrix& A,
                         const SVector& S, const UMatrix& U, const VMatrix& Vt) {
  using memory_space = typename AMatrix::memory_space;
  using Scalar       = typename AMatrix::non_const_value_type;
  using Magnitude    = typename SVector::non_const_value_type;
  using ALayout_t    = typename AMatrix::array_layout;
  using ULayout_t    = typename UMatrix::array_layout;
  using VLayout_t    = typename VMatrix::array_layout;

  static_assert(std::is_same_v<ALayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: A needs to have a Kokkos::LayoutLeft");
  static_assert(std::is_same_v<ULayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: U needs to have a Kokkos::LayoutLeft");
  static_assert(std::is_same_v<VLayout_t, Kokkos::LayoutLeft>,
                "KokkosLapack - svd: Vt needs to have a Kokkos::LayoutLeft");

  const rocblas_int m    = A.extent_int(0);
  const rocblas_int n    = A.extent_int(1);
  const rocblas_int lda  = A.stride(1);
  const rocblas_int ldu  = U.stride(1);
  const rocblas_int ldvt = Vt.stride(1);

  rocblas_svect UVecMode = rocblas_svect_all;
  if ((jobu[0] == 'S') || (jobu[0] == 's')) {
    UVecMode = rocblas_svect_singular;
  } else if ((jobu[0] == 'O') || (jobu[0] == 'o')) {
    UVecMode = rocblas_svect_overwrite;
  } else if ((jobu[0] == 'N') || (jobu[0] == 'n')) {
    UVecMode = rocblas_svect_none;
  }
  rocblas_svect VVecMode = rocblas_svect_all;
  if ((jobvt[0] == 'S') || (jobvt[0] == 's')) {
    VVecMode = rocblas_svect_singular;
  } else if ((jobvt[0] == 'O') || (jobvt[0] == 'o')) {
    VVecMode = rocblas_svect_overwrite;
  } else if ((jobvt[0] == 'N') || (jobvt[0] == 'n')) {
    VVecMode = rocblas_svect_none;
  }

  const rocblas_workmode WorkMode = rocblas_outofplace;

  Kokkos::View<rocblas_int, memory_space> info("svd info");
  Kokkos::View<Magnitude*, memory_space> rwork("svd rwork buffer", Kokkos::min(m, n) - 1);

  KokkosBlas::Impl::RocBlasSingleton& s = KokkosBlas::Impl::RocBlasSingleton::singleton();
  KOKKOSBLAS_IMPL_ROCBLAS_SAFE_CALL(rocblas_set_stream(s.handle, space.hip_stream()));
  if constexpr (std::is_same_v<Scalar, float>) {
    KOKKOSBLAS_IMPL_ROCBLAS_SAFE_CALL(rocsolver_sgesvd(s.handle, UVecMode, VVecMode, m, n, A.data(), lda, S.data(),
                                                       U.data(), ldu, Vt.data(), ldvt, rwork.data(), WorkMode,
                                                       info.data()));
  }
  if constexpr (std::is_same_v<Scalar, double>) {
    KOKKOSBLAS_IMPL_ROCBLAS_SAFE_CALL(rocsolver_dgesvd(s.handle, UVecMode, VVecMode, m, n, A.data(), lda, S.data(),
                                                       U.data(), ldu, Vt.data(), ldvt, rwork.data(), WorkMode,
                                                       info.data()));
  }
  if constexpr (std::is_same_v<Scalar, Kokkos::complex<float>>) {
    KOKKOSBLAS_IMPL_ROCBLAS_SAFE_CALL(rocsolver_cgesvd(
        s.handle, UVecMode, VVecMode, m, n, reinterpret_cast<rocblas_float_complex*>(A.data()), lda, S.data(),
        reinterpret_cast<rocblas_float_complex*>(U.data()), ldu, reinterpret_cast<rocblas_float_complex*>(Vt.data()),
        ldvt, rwork.data(), WorkMode, info.data()));
  }
  if constexpr (std::is_same_v<Scalar, Kokkos::complex<double>>) {
    KOKKOSBLAS_IMPL_ROCBLAS_SAFE_CALL(rocsolver_zgesvd(
        s.handle, UVecMode, VVecMode, m, n, reinterpret_cast<rocblas_double_complex*>(A.data()), lda, S.data(),
        reinterpret_cast<rocblas_double_complex*>(U.data()), ldu, reinterpret_cast<rocblas_double_complex*>(Vt.data()),
        ldvt, rwork.data(), WorkMode, info.data()));
  }
  KOKKOSBLAS_IMPL_ROCBLAS_SAFE_CALL(rocblas_set_stream(s.handle, NULL));
}

#define KOKKOSLAPACK_SVD_ROCSOLVER(SCALAR, LAYOUT, MEM_SPACE)                                                          \
  template <>                                                                                                          \
  struct SVD<                                                                                                          \
      Kokkos::HIP,                                                                                                     \
      Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::HIP, MEM_SPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
      Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                                  \
                   Kokkos::Device<Kokkos::HIP, MEM_SPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                   \
      Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::HIP, MEM_SPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
      Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::HIP, MEM_SPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>, \
      true,                                                                                                            \
      svd_eti_spec_avail<                                                                                              \
          Kokkos::HIP,                                                                                                 \
          Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::HIP, MEM_SPACE>,                                       \
                       Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                       \
          Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                              \
                       Kokkos::Device<Kokkos::HIP, MEM_SPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>,               \
          Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::HIP, MEM_SPACE>,                                       \
                       Kokkos::MemoryTraits<Kokkos::Unmanaged>>,                                                       \
          Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::HIP, MEM_SPACE>,                                       \
                       Kokkos::MemoryTraits<Kokkos::Unmanaged>>>::value> {                                             \
    using AMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::HIP, MEM_SPACE>,                             \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                             \
    using SVector = Kokkos::View<KokkosKernels::ArithTraits<SCALAR>::mag_type*, Kokkos::LayoutLeft,                    \
                                 Kokkos::Device<Kokkos::HIP, MEM_SPACE>, Kokkos::MemoryTraits<Kokkos::Unmanaged>>;     \
    using UMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::HIP, MEM_SPACE>,                             \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                             \
    using VMatrix = Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::HIP, MEM_SPACE>,                             \
                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;                                             \
                                                                                                                       \
    static void svd(const Kokkos::HIP& space, const char jobu[], const char jobvt[], const AMatrix& A,                 \
                    const SVector& S, const UMatrix& U, const VMatrix& Vt) {                                           \
      Kokkos::Profiling::pushRegion("KokkosLapack::svd[TPL_ROCSOLVER," #SCALAR "]");                                   \
      svd_print_specialization<Kokkos::HIP, AMatrix, SVector, UMatrix, VMatrix>();                                     \
                                                                                                                       \
      rocsolverSvdWrapper(space, jobu, jobvt, A, S, U, Vt);                                                            \
      Kokkos::Profiling::popRegion();                                                                                  \
    }                                                                                                                  \
  };

KOKKOSLAPACK_SVD_ROCSOLVER(float, Kokkos::LayoutLeft, Kokkos::HIPSpace)
KOKKOSLAPACK_SVD_ROCSOLVER(double, Kokkos::LayoutLeft, Kokkos::HIPSpace)
KOKKOSLAPACK_SVD_ROCSOLVER(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::HIPSpace)
KOKKOSLAPACK_SVD_ROCSOLVER(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::HIPSpace)

#if defined(KOKKOSKERNELS_INST_MEMSPACE_HIPMANAGEDSPACE)
KOKKOSLAPACK_SVD_ROCSOLVER(float, Kokkos::LayoutLeft, Kokkos::HIPManagedSpace)
KOKKOSLAPACK_SVD_ROCSOLVER(double, Kokkos::LayoutLeft, Kokkos::HIPManagedSpace)
KOKKOSLAPACK_SVD_ROCSOLVER(Kokkos::complex<float>, Kokkos::LayoutLeft, Kokkos::HIPManagedSpace)
KOKKOSLAPACK_SVD_ROCSOLVER(Kokkos::complex<double>, Kokkos::LayoutLeft, Kokkos::HIPManagedSpace)
#endif

}  // namespace Impl
}  // namespace KokkosLapack
#endif  // KOKKOSKERNELS_ENABLE_TPL_ROCSOLVER

#endif  // KOKKOSLAPACK_SVD_TPL_SPEC_DECL_HPP_
