ops/linear_solve.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/common/trace.hpp>
#include <uni20/core/scalar_concepts.hpp>
#include <uni20/linalg/backends/cpu/linear_solve.hpp>
#include <uni20/linalg/backends/lapack/linear_solve.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/linalg/operation_tags.hpp>
#include <uni20/linalg/precision.hpp>
#include <uni20/linalg/solve_info.hpp>
#include <uni20/tensor/concepts.hpp>
#include <uni20/tensor/copy.hpp>
#include <uni20/tensor/precision.hpp>
#include <uni20/tensor/tensor.hpp>
#if UNI20_HAS_MPLAPACK_FLOAT80
#include <uni20/linalg/backends/lapack/linear_solve_binary80.hpp>
#endif
 
#include <concepts>
#include <type_traits>
#include <utility>
 
namespace uni20::linalg
{
namespace detail
{
template <class CoefficientTensor, class RhsTensor>
concept CompatibleLinearSolveTensors =
    uni20::MutableRankedTensorView<CoefficientTensor, 2> && uni20::MutableRankedTensorView<RhsTensor, 2> &&
    uni20::RealOrComplex<uni20::tensor_element_t<CoefficientTensor>> &&
    std::same_as<uni20::tensor_element_t<CoefficientTensor>, uni20::tensor_element_t<RhsTensor>>;
 
template <uni20::RankedTensorView<2> CoefficientTensor, uni20::RankedTensorView<2> RhsTensor>
void require_linear_solve_shape(CoefficientTensor const& coefficients, RhsTensor const& right_hand_sides)
{
  ERROR_IF(coefficients.extent(0) != coefficients.extent(1), "solve requires a square coefficient matrix");
  ERROR_IF(coefficients.extent(0) != right_hand_sides.extent(0),
           "solve coefficient and right-hand-side row counts do not agree");
}
 
inline void require_solve_success(SolveInfo const& info)
{
  ERROR_IF(info.status == SolveStatus::singular, "found a singular matrix in solve");
  ERROR_IF(info.status == SolveStatus::nonfinite_input, "nonfinite input in solve");
  ERROR_IF(!info.succeeded(), "nonfinite result or rejected pivot in solve");
}
} // namespace detail
 
template <KernelBackendSelector BackendSelector, class A, class B>
  requires detail::CompatibleLinearSolveTensors<A, B>
[[nodiscard]] SolveInfo solve_inplace_with_info(BackendSelector&& selector, A&& a, B&& b, Precision p,
                                                SolveOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  require_precision<tensor_element_t<A>>(p);
  (void)p.bit_count();
  detail::require_linear_solve_shape(a, b);
  detail::require_solve_options(options);
  auto ad = mdspec_of(a);
  auto bd = mdspec_of(b);
  SolveInfo info;
  detail::dispatch_at_precision<tensor_element_t<A>>(std::forward<BackendSelector>(selector), linear_solve_op{}, p, ad,
                                                     bd, info, options);
  if (info.succeeded() && a.extent(0) != 0 && b.extent(1) != 0)
  {
    uni20::detail::record_output_precision(a, p);
    uni20::detail::record_output_precision(b, p);
  }
  return info;
}
 
template <KernelBackendSelector BackendSelector, class CoefficientTensor, class RhsTensor>
  requires detail::CompatibleLinearSolveTensors<CoefficientTensor, RhsTensor>
[[nodiscard]] SolveInfo solve_inplace_with_info(
    BackendSelector&& selector, CoefficientTensor&& coefficients, RhsTensor&& right_hand_sides,
    SolveOptions<uni20::make_real_t<uni20::tensor_element_t<CoefficientTensor>>> const& options = {})
{
  return solve_inplace_with_info(std::forward<BackendSelector>(selector), std::forward<CoefficientTensor>(coefficients),
                                 std::forward<RhsTensor>(right_hand_sides),
                                 common_default_precision(coefficients, right_hand_sides), options);
}
 
template <class CoefficientTensor, class RhsTensor>
  requires detail::CompatibleLinearSolveTensors<CoefficientTensor, RhsTensor>
[[nodiscard]] SolveInfo solve_inplace_with_info(
    CoefficientTensor&& coefficients, RhsTensor&& right_hand_sides,
    SolveOptions<uni20::make_real_t<uni20::tensor_element_t<CoefficientTensor>>> const& options = {})
{
  auto selector = select_backend(linear_solve_op{}, coefficients, right_hand_sides);
  return solve_inplace_with_info(selector, std::forward<CoefficientTensor>(coefficients),
                                 std::forward<RhsTensor>(right_hand_sides), options);
}
 
template <KernelBackendSelector BackendSelector, class CoefficientTensor, class RhsTensor>
  requires detail::CompatibleLinearSolveTensors<CoefficientTensor, RhsTensor>
void solve_inplace(BackendSelector&& selector, CoefficientTensor&& coefficients, RhsTensor&& right_hand_sides,
                   std::optional<Precision> requested = {})
{
  auto const info =
      solve_inplace_with_info(std::forward<BackendSelector>(selector), std::forward<CoefficientTensor>(coefficients),
                              std::forward<RhsTensor>(right_hand_sides),
                              requested ? *requested : common_default_precision(coefficients, right_hand_sides));
  detail::require_solve_success(info);
}
 
template <class CoefficientTensor, class RhsTensor>
  requires detail::CompatibleLinearSolveTensors<CoefficientTensor, RhsTensor>
void solve_inplace(CoefficientTensor&& coefficients, RhsTensor&& right_hand_sides, std::optional<Precision> p = {})
{
  detail::require_linear_solve_shape(coefficients, right_hand_sides);
  auto selector = select_backend(linear_solve_op{}, coefficients, right_hand_sides);
  solve_inplace(selector, std::forward<CoefficientTensor>(coefficients), std::forward<RhsTensor>(right_hand_sides), p);
}
 
template <KernelBackendSelector BackendSelector, uni20::RankedTensorView<2> CoefficientTensor,
          uni20::RankedTensorView<2> RhsTensor>
  requires uni20::RealOrComplex<uni20::tensor_element_t<CoefficientTensor>> &&
           std::same_as<uni20::tensor_element_t<CoefficientTensor>, uni20::tensor_element_t<RhsTensor>>
[[nodiscard]] auto solve(BackendSelector&& selector, CoefficientTensor const& coefficients,
                         RhsTensor const& right_hand_sides, std::optional<Precision> requested = {})
{
  auto p = requested ? *requested : common_default_precision(coefficients, right_hand_sides);
  require_precision<tensor_element_t<CoefficientTensor>>(p);
  detail::require_linear_solve_shape(coefficients, right_hand_sides);
  auto coefficient_work = uni20::make_tensor<uni20::ColumnMajor>(coefficients);
  auto solution = uni20::make_tensor<uni20::ColumnMajor>(right_hand_sides);
  solve_inplace(std::forward<BackendSelector>(selector), coefficient_work, solution, p);
  uni20::detail::record_output_precision(solution, p);
  return solution;
}
 
template <uni20::RankedTensorView<2> CoefficientTensor, uni20::RankedTensorView<2> RhsTensor>
  requires uni20::RealOrComplex<uni20::tensor_element_t<CoefficientTensor>> &&
           std::same_as<uni20::tensor_element_t<CoefficientTensor>, uni20::tensor_element_t<RhsTensor>>
[[nodiscard]] auto solve(CoefficientTensor const& coefficients, RhsTensor const& right_hand_sides,
                         std::optional<Precision> requested = {})
{
  using work_type = decltype(uni20::make_tensor<uni20::ColumnMajor>(right_hand_sides));
  auto selector = select_backend_for<work_type, work_type>(linear_solve_op{});
  return solve(selector, coefficients, right_hand_sides, requested);
}
 
template <class A, class B>
  requires detail::CompatibleLinearSolveTensors<A, B>
[[nodiscard]] SolveInfo solve_inplace_with_info(A&& a, B&& b, Precision p,
                                                SolveOptions<make_real_t<tensor_element_t<A>>> const& options = {})
{
  auto selector = select_backend(linear_solve_op{}, a, b);
  return solve_inplace_with_info(selector, std::forward<A>(a), std::forward<B>(b), p, options);
}
 
} // namespace uni20::linalg