ops/self_adjoint_eigh.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/common/trace.hpp>
#include <uni20/linalg/backends/lapack/self_adjoint_eigh.hpp>
#include <uni20/linalg/blas/mdspan_matrix.hpp>
#include <uni20/linalg/decomposition_precision.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/linalg/operation_tags.hpp>
#include <uni20/tensor/conjugate_inplace.hpp>
#include <uni20/tensor/copy.hpp>
#include <uni20/tensor/output.hpp>
#include <uni20/tensor/tensor.hpp>
 
#include <array>
#include <concepts>
#include <type_traits>
#include <utility>
 
namespace uni20::linalg
{
namespace detail
{
using self_adjoint_eigenvalue_extents = stdex::dextents<uni20::index_type, 1>;
 
template <uni20::MutableRankedTensorView<1> EigenvalueTensor, uni20::MutableRankedTensorView<2> MatrixTensor>
[[nodiscard]] self_adjoint_eigh_op prepare_self_adjoint_eigh(EigenvalueTensor& eigenvalues, MatrixTensor& matrix_work,
                                                             SelfAdjointEighOptions options)
{
  ERROR_IF(matrix_work.extent(0) != matrix_work.extent(1), "self-adjoint eigensystem requires a square matrix");
  uni20::prepare_output(eigenvalues,
                        self_adjoint_eigenvalue_extents{static_cast<uni20::index_type>(matrix_work.extent(0))});
  return {.compute_vectors = options.compute_vectors, .triangle = options.triangle};
}
 
template <KernelBackendSelector BackendSelector, uni20::MutableRankedTensorView<1> EigenvalueTensor,
          uni20::MutableRankedTensorView<2> MatrixTensor>
void dispatch_self_adjoint_eigh(BackendSelector&& selector, self_adjoint_eigh_op operation,
                                EigenvalueTensor& eigenvalues, MatrixTensor& matrix_work, Precision p)
{
  auto eigenvalue_descriptor = uni20::mdspec_of(eigenvalues);
  auto matrix_descriptor = uni20::mdspec_of(matrix_work);
  detail::dispatch_at_precision<uni20::tensor_element_t<MatrixTensor>>(
      std::forward<BackendSelector>(selector), operation, p, eigenvalue_descriptor, matrix_descriptor);
 
  uni20::detail::record_output_precision(eigenvalues, p);
  uni20::detail::record_output_precision(matrix_work, p);
}
} // namespace detail
 
template <KernelBackendSelector BackendSelector, uni20::MutableRankedTensorView<1> EigenvalueTensor,
          uni20::MutableRankedTensorView<2> MatrixTensor>
void self_adjoint_eigh(BackendSelector&& selector, EigenvalueTensor&& eigenvalues, MatrixTensor&& matrix_work,
                       SelfAdjointEighOptions options = {}, std::optional<Precision> requested = {})
{
  auto const p = detail::decomposition_precision(matrix_work, requested);
  auto operation = detail::prepare_self_adjoint_eigh(eigenvalues, matrix_work, options);
  detail::dispatch_self_adjoint_eigh(std::forward<BackendSelector>(selector), operation, eigenvalues, matrix_work, p);
}
 
template <uni20::MutableRankedTensorView<1> EigenvalueTensor, uni20::MutableRankedTensorView<2> MatrixTensor>
void self_adjoint_eigh(EigenvalueTensor&& eigenvalues, MatrixTensor&& matrix_work, SelfAdjointEighOptions options = {},
                       std::optional<Precision> requested = {})
{
  auto const p = detail::decomposition_precision(matrix_work, requested);
  auto operation = detail::prepare_self_adjoint_eigh(eigenvalues, matrix_work, options);
  auto selector = select_backend(operation, eigenvalues, matrix_work);
  detail::dispatch_self_adjoint_eigh(selector, operation, eigenvalues, matrix_work, p);
}
 
template <class EigenvalueTensor, class EigenvectorTensor> struct SelfAdjointEighResult
{
    using eigenvalue_tensor_type = EigenvalueTensor;
    using eigenvector_tensor_type = EigenvectorTensor;
 
    EigenvalueTensor eigenvalues;
    EigenvectorTensor eigenvectors;
};
 
namespace detail
{
template <class MatrixTensor>
using self_adjoint_eigh_reuse_matrix_t =
    std::conditional_t<std::same_as<typename std::remove_cvref_t<MatrixTensor>::layout_type, uni20::ColumnMajor>,
                       std::remove_cvref_t<MatrixTensor>,
                       typename std::remove_cvref_t<MatrixTensor>::template rebind_layout_type<stdex::layout_stride>>;
 
template <class MatrixTensor> consteval bool can_transfer_self_adjoint_eigh_storage()
{
  using matrix_type = std::remove_cvref_t<MatrixTensor>;
  if constexpr (requires {
                  typename matrix_type::storage_policy;
                  typename matrix_type::layout_type;
                  typename matrix_type::accessor_factory_type;
                  typename matrix_type::storage_type;
                  typename matrix_type::template rebind_layout_type<stdex::layout_stride>;
                })
  {
    using work_type = self_adjoint_eigh_reuse_matrix_t<matrix_type>;
    return std::same_as<typename matrix_type::storage_policy, uni20::HostStorage> &&
           std::same_as<typename matrix_type::accessor_factory_type, uni20::DefaultAccessorFactory> &&
           uni20::DefaultAccessorMdspanLike<uni20::mutable_immediate_tensor_mdspan_t<matrix_type>> &&
           requires(matrix_type& matrix, typename work_type::mapping_type mapping,
                    typename work_type::storage_type storage) {
             { std::move(matrix).release_storage() } -> std::same_as<typename work_type::storage_type>;
             { work_type::adopt_storage(std::move(mapping), std::move(storage)) } -> std::same_as<work_type>;
           };
  }
  else
  {
    return false;
  }
}
 
constexpr MatrixTriangle transposed_triangle(MatrixTriangle triangle)
{
  switch (triangle)
  {
    case MatrixTriangle::Upper:
      return MatrixTriangle::Lower;
    case MatrixTriangle::Lower:
      return MatrixTriangle::Upper;
  }
  PANIC("invalid MatrixTriangle", std::to_underlying(triangle));
}
 
template <class WorkMatrix, uni20::RankedTensorView<2> MatrixTensor>
[[nodiscard]] WorkMatrix materialize_self_adjoint_eigh_work_matrix(MatrixTensor const& matrix)
{
  WorkMatrix result = [&] {
    if constexpr (std::same_as<typename WorkMatrix::layout_type, uni20::ColumnMajor>)
    {
      return WorkMatrix(matrix.extents());
    }
    else
    {
      return WorkMatrix(matrix.extents(), uni20::layout::LayoutLeft{});
    }
  }();
  uni20::copy(result, matrix);
  return result;
}
 
template <class WorkMatrix, uni20::RankedImmediateTensorView<2> MatrixTensor>
[[nodiscard]] auto column_major_reuse_mapping(MatrixTensor const& matrix, uni20::blas_int leading_dimension) ->
    typename WorkMatrix::mapping_type
{
  if constexpr (std::same_as<typename WorkMatrix::layout_type, uni20::ColumnMajor>)
  {
    return typename WorkMatrix::mapping_type(matrix.extents());
  }
  else
  {
    using index_type = typename WorkMatrix::index_type;
    return typename WorkMatrix::mapping_type(
        matrix.extents(), std::array<index_type, 2>{index_type{1}, static_cast<index_type>(leading_dimension)});
  }
}
 
template <class MatrixTensor>
  requires uni20::RankedTensorView<MatrixTensor, 2>
[[nodiscard]] constexpr auto select_self_adjoint_eigh_backend(MatrixTriangle triangle)
{
  using real_type = uni20::make_real_t<uni20::tensor_element_t<MatrixTensor>>;
  using eigenvalue_tensor = uni20::Tensor<real_type, 1>;
  return select_backend_for<eigenvalue_tensor, std::remove_cvref_t<MatrixTensor>>(
      self_adjoint_eigh_op{.compute_vectors = true, .triangle = triangle});
}
 
template <KernelBackendSelector BackendSelector, uni20::OwningTensor MatrixTensor>
  requires uni20::MutableRankedImmediateTensorView<MatrixTensor, 2>
[[nodiscard]] auto eigh_from_work_matrix(BackendSelector&& selector, MatrixTensor matrix_work, MatrixTriangle triangle,
                                         Precision p)
{
  using real_type = uni20::make_real_t<uni20::tensor_element_t<MatrixTensor>>;
  uni20::Tensor<real_type, 1> eigenvalues;
  self_adjoint_eigh(std::forward<BackendSelector>(selector), eigenvalues, matrix_work,
                    SelfAdjointEighOptions{.compute_vectors = true, .triangle = triangle}, p);
  return SelfAdjointEighResult<decltype(eigenvalues), MatrixTensor>{.eigenvalues = std::move(eigenvalues),
                                                                    .eigenvectors = std::move(matrix_work)};
}
 
template <KernelBackendSelector BackendSelector, uni20::OwningTensor MatrixTensor>
  requires uni20::MutableRankedImmediateTensorView<MatrixTensor, 2>
[[nodiscard]] auto eigh_from_consumed_matrix(BackendSelector&& selector, MatrixTensor&& matrix, MatrixTriangle triangle,
                                             Precision p)
{
  using matrix_type = std::remove_cvref_t<MatrixTensor>;
  using work_type = self_adjoint_eigh_reuse_matrix_t<matrix_type>;
 
  auto const matrix_span = matrix.mdspan();
  auto const stage = uni20::linalg::blas::try_mdspan_matrix_stage(matrix_span);
  if (stage && !stage->needs_conjugation)
  {
    bool const transpose_storage = stage->unit_stride_axis == 1;
    MatrixTriangle const work_triangle = transpose_storage ? transposed_triangle(triangle) : triangle;
    auto work_mapping = column_major_reuse_mapping<work_type>(matrix, stage->nonunit_stride);
    auto storage = std::move(matrix).release_storage();
    auto work_matrix = work_type::adopt_storage(std::move(work_mapping), std::move(storage));
    auto result =
        eigh_from_work_matrix(std::forward<BackendSelector>(selector), std::move(work_matrix), work_triangle, p);
    if constexpr (uni20::Complex<uni20::tensor_element_t<matrix_type>>)
    {
      if (transpose_storage) uni20::conjugate_inplace(result.eigenvectors);
    }
    return result;
  }
 
  auto work_matrix = materialize_self_adjoint_eigh_work_matrix<work_type>(matrix);
  return eigh_from_work_matrix(std::forward<BackendSelector>(selector), std::move(work_matrix), triangle, p);
}
} // namespace detail
 
template <KernelBackendSelector BackendSelector, uni20::RankedTensorView<2> MatrixTensor>
[[nodiscard]] auto eigh(BackendSelector&& selector, MatrixTensor const& matrix,
                        MatrixTriangle triangle = MatrixTriangle::Upper, std::optional<Precision> requested = {})
{
  auto const p = detail::decomposition_precision(matrix, requested);
  ERROR_IF(matrix.extent(0) != matrix.extent(1), "self-adjoint eigensystem requires a square matrix");
  auto matrix_work = uni20::make_tensor<uni20::ColumnMajor>(matrix);
  return detail::eigh_from_work_matrix(std::forward<BackendSelector>(selector), std::move(matrix_work), triangle, p);
}
 
template <uni20::RankedTensorView<2> MatrixTensor>
[[nodiscard]] auto eigh(MatrixTensor const& matrix, MatrixTriangle triangle = MatrixTriangle::Upper,
                        std::optional<Precision> requested = {})
{
  auto const p = detail::decomposition_precision(matrix, requested);
  ERROR_IF(matrix.extent(0) != matrix.extent(1), "self-adjoint eigensystem requires a square matrix");
  auto matrix_work = uni20::make_tensor<uni20::ColumnMajor>(matrix);
  auto selector = detail::select_self_adjoint_eigh_backend<decltype(matrix_work)>(triangle);
  return detail::eigh_from_work_matrix(std::move(selector), std::move(matrix_work), triangle, p);
}
 
template <KernelBackendSelector BackendSelector, class MatrixTensor>
  requires uni20::OwningTensor<MatrixTensor> && uni20::MutableRankedImmediateTensorView<MatrixTensor, 2> &&
           (!std::is_lvalue_reference_v<MatrixTensor>) && (!std::is_const_v<std::remove_reference_t<MatrixTensor>>)
[[nodiscard]] auto eigh(BackendSelector&& selector, MatrixTensor&& matrix,
                        MatrixTriangle triangle = MatrixTriangle::Upper, std::optional<Precision> requested = {})
{
  auto const p = detail::decomposition_precision(matrix, requested);
  ERROR_IF(matrix.extent(0) != matrix.extent(1), "self-adjoint eigensystem requires a square matrix");
  if constexpr (detail::can_transfer_self_adjoint_eigh_storage<MatrixTensor>())
  {
    return detail::eigh_from_consumed_matrix(std::forward<BackendSelector>(selector),
                                             std::forward<MatrixTensor>(matrix), triangle, p);
  }
  else
  {
    auto matrix_work = uni20::make_tensor<uni20::ColumnMajor>(matrix);
    return detail::eigh_from_work_matrix(std::forward<BackendSelector>(selector), std::move(matrix_work), triangle, p);
  }
}
 
template <class MatrixTensor>
  requires uni20::OwningTensor<MatrixTensor> && uni20::MutableRankedImmediateTensorView<MatrixTensor, 2> &&
           (!std::is_lvalue_reference_v<MatrixTensor>) && (!std::is_const_v<std::remove_reference_t<MatrixTensor>>)
[[nodiscard]] auto eigh(MatrixTensor&& matrix, MatrixTriangle triangle = MatrixTriangle::Upper,
                        std::optional<Precision> requested = {})
{
  auto const p = detail::decomposition_precision(matrix, requested);
  ERROR_IF(matrix.extent(0) != matrix.extent(1), "self-adjoint eigensystem requires a square matrix");
  if constexpr (detail::can_transfer_self_adjoint_eigh_storage<MatrixTensor>())
  {
    using work_type = detail::self_adjoint_eigh_reuse_matrix_t<std::remove_cvref_t<MatrixTensor>>;
    auto selector = detail::select_self_adjoint_eigh_backend<work_type>(triangle);
    return detail::eigh_from_consumed_matrix(std::move(selector), std::forward<MatrixTensor>(matrix), triangle, p);
  }
  else
  {
    auto matrix_work = uni20::make_tensor<uni20::ColumnMajor>(matrix);
    auto selector = detail::select_self_adjoint_eigh_backend<decltype(matrix_work)>(triangle);
    return detail::eigh_from_work_matrix(std::move(selector), std::move(matrix_work), triangle, p);
  }
}
 
template <RankedTensorView<2> Matrix>
[[nodiscard]] auto eigh(Matrix&& matrix, Precision p, MatrixTriangle triangle = MatrixTriangle::Upper)
{
  return uni20::linalg::eigh(std::forward<Matrix>(matrix), triangle, std::optional<Precision>{p});
}
 
template <KernelBackendSelector Selector, RankedTensorView<2> Matrix>
[[nodiscard]] auto eigh(Selector&& selector, Matrix&& matrix, Precision p,
                        MatrixTriangle triangle = MatrixTriangle::Upper)
{
  return uni20::linalg::eigh(std::forward<Selector>(selector), std::forward<Matrix>(matrix), triangle,
                             std::optional<Precision>{p});
}
 
template <MutableRankedTensorView<1> V, MutableRankedTensorView<2> A>
void self_adjoint_eigh(V&& eigenvalues, A&& matrix_work, Precision p, SelfAdjointEighOptions options = {})
{
  uni20::linalg::self_adjoint_eigh(std::forward<V>(eigenvalues), std::forward<A>(matrix_work), options,
                                   std::optional<Precision>{p});
}
 
template <KernelBackendSelector Selector, MutableRankedTensorView<1> V, MutableRankedTensorView<2> A>
void self_adjoint_eigh(Selector&& selector, V&& eigenvalues, A&& matrix_work, Precision p,
                       SelfAdjointEighOptions options = {})
{
  uni20::linalg::self_adjoint_eigh(std::forward<Selector>(selector), std::forward<V>(eigenvalues),
                                   std::forward<A>(matrix_work), options, std::optional<Precision>{p});
}
 
} // namespace uni20::linalg