backends/lapack/self_adjoint_eigh.hpp

Go to the documentation of this file.

#pragma once
 
#include "common.hpp"
 
#include <uni20/backend/blas/blas_int.hpp>
#include <uni20/backend/lapack/lapack.hpp>
#include <uni20/common/trace.hpp>
#include <uni20/core/scalar_concepts.hpp>
#include <uni20/linalg/blas/mdspan_matrix.hpp>
#include <uni20/linalg/blas/mdspan_vector.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/linalg/operation_tags.hpp>
#include <uni20/tensor/access.hpp>
#include <uni20/tensor/concepts.hpp>
 
#include <algorithm>
#include <concepts>
#include <cstddef>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace uni20::linalg
{
namespace lapack_detail
{
inline char lapack_triangle(MatrixTriangle triangle)
{
  switch (triangle)
  {
    case MatrixTriangle::Upper:
      return 'U';
    case MatrixTriangle::Lower:
      return 'L';
  }
  PANIC("invalid MatrixTriangle", std::to_underlying(triangle));
}
 
template <uni20::MutableRankedStridedMdspanLike<1> EigenvalueMdspan,
          uni20::MutableRankedStridedMdspanLike<2> MatrixMdspan>
consteval auto self_adjoint_eigh_acceptance()
{
  using matrix_scalar = std::remove_cv_t<typename MatrixMdspan::element_type>;
  using eigenvalue_scalar = std::remove_cv_t<typename EigenvalueMdspan::element_type>;
  if constexpr (uni20::LapackScalar<matrix_scalar> && uni20::LapackReal<eigenvalue_scalar> &&
                std::same_as<eigenvalue_scalar, uni20::make_real_t<matrix_scalar>> &&
                uni20::DefaultAccessorMdspanLike<EigenvalueMdspan> && uni20::DefaultAccessorMdspanLike<MatrixMdspan>)
  {
    return kernel_types_maybe;
  }
  else
  {
    return kernel_types_no;
  }
}
 
template <uni20::MutableRankedStridedMdspanLike<1> EigenvalueMdspan,
          uni20::MutableRankedStridedMdspanLike<2> MatrixMdspan>
KernelAttempt try_self_adjoint_eigh(self_adjoint_eigh_op const& op, EigenvalueMdspan& eigenvalues,
                                    MatrixMdspan& matrix_work)
{
  using matrix_type = std::remove_cvref_t<MatrixMdspan>;
  using matrix_scalar = std::remove_cv_t<typename matrix_type::element_type>;
  using real_type = uni20::make_real_t<matrix_scalar>;
 
  CHECK_EQUAL(matrix_work.extent(0), matrix_work.extent(1));
  std::size_t const n = static_cast<std::size_t>(matrix_work.extent(0));
  CHECK_EQUAL(static_cast<std::size_t>(eigenvalues.extent(0)), n);
 
  blas_int const order = uni20::blas::try_blas_int(n);
  if (!uni20::blas::is_valid_blas_int(order))
  {
    return KernelAttempt::unsupported_shape;
  }
  if (order == 0)
  {
    return KernelAttempt::success;
  }
 
  auto matrix = uni20::linalg::blas::try_lapack_writable_matrix(matrix_work);
  auto values = uni20::linalg::blas::try_lapack_writable_vector(eigenvalues);
  if (!matrix || !values)
  {
    return KernelAttempt::unsupported_layout;
  }
  CHECK_EQUAL(matrix->rows, order);
  CHECK_EQUAL(matrix->cols, order);
  CHECK_EQUAL(values->size, order);
 
  char const jobz = op.compute_vectors ? 'V' : 'N';
  char const uplo = lapack_detail::lapack_triangle(op.triangle);
  if constexpr (uni20::LapackReal<matrix_scalar>)
  {
    matrix_scalar work_query{};
    uni20::lapack::syev(jobz, uplo, order, matrix->data, matrix->leading_dimension, values->data, &work_query, -1);
    blas_int const lwork = lapack_detail::workspace_size(work_query);
    std::vector<matrix_scalar> work(static_cast<std::size_t>(lwork), matrix_scalar{});
    uni20::lapack::syev(jobz, uplo, order, matrix->data, matrix->leading_dimension, values->data, work.data(), lwork);
  }
  else
  {
    auto const three_order = lapack_detail::try_size_product(3, n);
    if (!three_order)
    {
      return KernelAttempt::unsupported_shape;
    }
    std::size_t const rwork_size = std::max<std::size_t>(1, *three_order - 2);
    std::vector<real_type> rwork(rwork_size, real_type{});
 
    matrix_scalar work_query{};
    uni20::lapack::heev(jobz, uplo, order, matrix->data, matrix->leading_dimension, values->data, &work_query, -1,
                        rwork.data());
    blas_int const lwork = lapack_detail::workspace_size(work_query);
    std::vector<matrix_scalar> work(static_cast<std::size_t>(lwork), matrix_scalar{});
    uni20::lapack::heev(jobz, uplo, order, matrix->data, matrix->leading_dimension, values->data, work.data(), lwork,
                        rwork.data());
  }
 
  return KernelAttempt::success;
}
} // namespace lapack_detail
 
template <uni20::MutableRankedStridedMdspecLike<1> EigenvalueMdspan,
          uni20::MutableRankedStridedMdspecLike<2> MatrixMdspan>
  requires LapackScalar<typename MatrixMdspan::value_type> && uni20::HostWritableMdspec<EigenvalueMdspan> &&
           uni20::HostWritableMdspec<MatrixMdspan>
consteval auto kernel_accepts_types(LapackBackend const&, self_adjoint_eigh_op const&, EigenvalueMdspan&, MatrixMdspan&)
{
  using eigenvalue_span = uni20::host_write_mdspan_t<EigenvalueMdspan>;
  using matrix_span = uni20::host_write_mdspan_t<MatrixMdspan>;
  constexpr auto acceptance = lapack_detail::self_adjoint_eigh_acceptance<eigenvalue_span, matrix_span>();
  if constexpr (acceptance == KernelTypeAcceptance::no)
    return kernel_types_no;
  else
    return kernel_types_maybe;
}
 
template <uni20::MutableRankedStridedMdspecLike<1> EigenvalueMdspan,
          uni20::MutableRankedStridedMdspecLike<2> MatrixMdspan>
  requires LapackScalar<typename MatrixMdspan::value_type> && uni20::HostWritableMdspec<EigenvalueMdspan> &&
           uni20::HostWritableMdspec<MatrixMdspan>
KernelAttempt try_kernel(LapackBackend, self_adjoint_eigh_op const& operation, EigenvalueMdspan& eigenvalues,
                         MatrixMdspan& matrix_work)
{
  return lapack_detail::with_host_write_mdspans(
      [&](auto& eigenvalue_span, auto& matrix_span) {
        return lapack_detail::try_self_adjoint_eigh(operation, eigenvalue_span, matrix_span);
      },
      eigenvalues, matrix_work);
}
 
} // namespace uni20::linalg
 
#if UNI20_HAS_MPLAPACK_MPREAL || UNI20_HAS_MPLAPACK_FLOAT80
#include "decompositions_packed.hpp"
#endif