backends/cpu/linear_solve.hpp

Go to the documentation of this file.

#pragma once
 
#include "lu_common.hpp"
#include <uni20/common/trace.hpp>
#include <uni20/core/scalar_concepts.hpp>
#include <uni20/core/scalar_traits.hpp>
#include <uni20/linalg/backends/linear_solve_common.hpp>
#include <uni20/linalg/backends/packed_matrix.hpp>
#include <uni20/linalg/dispatch.hpp>
#include <uni20/linalg/operation_tags.hpp>
#include <uni20/mdspan/concepts.hpp>
#include <uni20/tensor/access.hpp>
#include <uni20/tensor/concepts.hpp>
 
#include <cmath>
#include <concepts>
#include <cstddef>
#include <type_traits>
#include <utility>
 
namespace uni20::linalg
{
namespace detail::cpu_reference
{
 
template <class CoefficientMdspan, class RhsMdspan> consteval auto linear_solve_acceptance()
{
  using coefficient_type = std::remove_cvref_t<CoefficientMdspan>;
  using rhs_type = std::remove_cvref_t<RhsMdspan>;
  using scalar_type = typename coefficient_type::value_type;
  using coefficient_index = typename coefficient_type::index_type;
  using rhs_index = typename rhs_type::index_type;
 
  if constexpr (!uni20::has_runtime_precision_v<scalar_type> && uni20::RealOrComplex<scalar_type> &&
                std::same_as<scalar_type, typename rhs_type::value_type> &&
                requires(CoefficientMdspan& coefficients, RhsMdspan& rhs, coefficient_index coefficient_i,
                         rhs_index rhs_i, scalar_type value) {
                  static_cast<scalar_type>(coefficients[coefficient_i, coefficient_i]);
                  static_cast<scalar_type>(rhs[rhs_i, rhs_i]);
                  coefficients[coefficient_i, coefficient_i] = value;
                  rhs[rhs_i, rhs_i] = value;
                  value / value;
                  value -= value * value;
                })
    return kernel_types_yes;
  else
    return kernel_types_no;
}
 
template <class CoefficientMdspan, class RhsMdspan>
KernelAttempt linear_solve(CoefficientMdspan& coefficients, RhsMdspan& right_hand_sides, SolveInfo& info,
                           SolveOptions<uni20::make_real_t<typename CoefficientMdspan::value_type>> const& options)
{
  using coefficient_type = std::remove_cvref_t<CoefficientMdspan>;
  using rhs_type = std::remove_cvref_t<RhsMdspan>;
  using scalar_type = typename coefficient_type::value_type;
  using real_type = uni20::make_real_t<scalar_type>;
  using coefficient_index = typename coefficient_type::index_type;
  using rhs_index = typename rhs_type::index_type;
 
  CHECK_EQUAL(coefficients.extent(0), coefficients.extent(1));
  CHECK_EQUAL(coefficients.extent(0), right_hand_sides.extent(0));
  std::size_t const order = static_cast<std::size_t>(coefficients.extent(0));
  std::size_t const rhs_count = static_cast<std::size_t>(right_hand_sides.extent(1));
  // Direct mdspan dispatch bypasses the Tensor frontend's option validation.
  require_solve_options(options);
  info = {};
  if (order == 0 || rhs_count == 0) return KernelAttempt::success;
  real_type scale{};
  if (!prepare_linear_solve(coefficients, right_hand_sides, scale, info)) return KernelAttempt::success;
 
  using std::abs;
  for (std::size_t k = 0; k < order; ++k)
  {
    auto const coefficient_k = static_cast<coefficient_index>(k);
    std::size_t pivot_row = k;
    real_type pivot_value = abs(static_cast<scalar_type>(coefficients[coefficient_k, coefficient_k]));
    for (std::size_t row = k + 1; row < order; ++row)
    {
      auto const coefficient_row = static_cast<coefficient_index>(row);
      real_type const candidate = abs(static_cast<scalar_type>(coefficients[coefficient_row, coefficient_k]));
      if (candidate > pivot_value)
      {
        pivot_value = candidate;
        pivot_row = row;
      }
    }
 
    if (!check_solve_pivot(pivot_value, scale, options, k, info)) return KernelAttempt::success;
    swap_rows(coefficients, k, pivot_row, order);
    swap_rows(right_hand_sides, k, pivot_row, rhs_count);
 
    scalar_type const pivot = static_cast<scalar_type>(coefficients[coefficient_k, coefficient_k]);
    for (std::size_t row = k + 1; row < order; ++row)
    {
      auto const coefficient_row = static_cast<coefficient_index>(row);
      scalar_type const factor = static_cast<scalar_type>(coefficients[coefficient_row, coefficient_k]) / pivot;
      if (factor == scalar_type{}) continue;
 
      coefficients[coefficient_row, coefficient_k] = scalar_type{};
      for (std::size_t col = k + 1; col < order; ++col)
      {
        auto const coefficient_col = static_cast<coefficient_index>(col);
        scalar_type value = static_cast<scalar_type>(coefficients[coefficient_row, coefficient_col]);
        value -= factor * static_cast<scalar_type>(coefficients[coefficient_k, coefficient_col]);
        if (!uni20::isfinite(value))
        {
          info.status = SolveStatus::nonfinite_result;
          return KernelAttempt::success;
        }
        coefficients[coefficient_row, coefficient_col] = value;
      }
      for (std::size_t col = 0; col < rhs_count; ++col)
      {
        auto const rhs_row = static_cast<rhs_index>(row);
        auto const rhs_k = static_cast<rhs_index>(k);
        auto const rhs_col = static_cast<rhs_index>(col);
        scalar_type value = static_cast<scalar_type>(right_hand_sides[rhs_row, rhs_col]);
        value -= factor * static_cast<scalar_type>(right_hand_sides[rhs_k, rhs_col]);
        if (!uni20::isfinite(value))
        {
          info.status = SolveStatus::nonfinite_result;
          return KernelAttempt::success;
        }
        right_hand_sides[rhs_row, rhs_col] = value;
      }
    }
  }
 
  for (std::size_t row = order; row-- > 0;)
  {
    auto const coefficient_row = static_cast<coefficient_index>(row);
    scalar_type const pivot = static_cast<scalar_type>(coefficients[coefficient_row, coefficient_row]);
    for (std::size_t col = 0; col < rhs_count; ++col)
    {
      auto const rhs_row = static_cast<rhs_index>(row);
      auto const rhs_col = static_cast<rhs_index>(col);
      scalar_type value = static_cast<scalar_type>(right_hand_sides[rhs_row, rhs_col]);
      for (std::size_t k = row + 1; k < order; ++k)
      {
        auto const coefficient_k = static_cast<coefficient_index>(k);
        auto const rhs_k = static_cast<rhs_index>(k);
        value -= static_cast<scalar_type>(coefficients[coefficient_row, coefficient_k]) *
                 static_cast<scalar_type>(right_hand_sides[rhs_k, rhs_col]);
      }
      value /= pivot;
      if (!uni20::isfinite(value))
      {
        info.status = SolveStatus::nonfinite_result;
        return KernelAttempt::success;
      }
      right_hand_sides[rhs_row, rhs_col] = value;
    }
  }
 
  return KernelAttempt::success;
}
 
} // namespace detail::cpu_reference
 
template <uni20::MutableRankedMdspecLike<2> CoefficientMdspec, uni20::MutableRankedMdspecLike<2> RhsMdspec>
  requires uni20::HostWritableMdspec<CoefficientMdspec> && uni20::HostWritableMdspec<RhsMdspec>
consteval auto kernel_accepts_types(CpuReferenceBackend const&, linear_solve_op const&, CoefficientMdspec&, RhsMdspec&,
                                    SolveInfo&,
                                    SolveOptions<uni20::make_real_t<typename CoefficientMdspec::value_type>> const&)
{
  using coefficient_span = uni20::host_write_mdspan_t<CoefficientMdspec>;
  using rhs_span = uni20::host_write_mdspan_t<RhsMdspec>;
  constexpr auto acceptance = detail::cpu_reference::linear_solve_acceptance<coefficient_span, rhs_span>();
  if constexpr (acceptance == KernelTypeAcceptance::yes)
    return kernel_types_yes;
  else
    return kernel_types_no;
}
 
template <uni20::MutableRankedMdspecLike<2> CoefficientMdspec, uni20::MutableRankedMdspecLike<2> RhsMdspec>
  requires uni20::HostWritableMdspec<CoefficientMdspec> && uni20::HostWritableMdspec<RhsMdspec>
KernelAttempt try_kernel(CpuReferenceBackend, linear_solve_op const&, CoefficientMdspec& coefficients,
                         RhsMdspec& right_hand_sides, SolveInfo& info,
                         SolveOptions<uni20::make_real_t<typename CoefficientMdspec::value_type>> const& options)
{
  auto coefficient_access = acquire_host_write_access_sync(coefficients);
  auto rhs_access = acquire_host_write_access_sync(right_hand_sides);
  auto coefficient_span = coefficient_access.mdspan();
  auto rhs_span = rhs_access.mdspan();
  return detail::cpu_reference::linear_solve(coefficient_span, rhs_span, info, options);
}
 
template <MutableRankedMdspecLike<2> A, MutableRankedMdspecLike<2> B>
  requires HostWritableMdspec<A> && HostWritableMdspec<B> && RealOrComplex<typename A::value_type> &&
           has_runtime_precision_v<typename A::value_type> &&
           std::same_as<typename A::value_type, typename B::value_type>
consteval auto kernel_accepts_types(CpuReferenceBackend, linear_solve_op const&, A&, B&, SolveInfo&,
                                    SolveOptions<make_real_t<typename A::value_type>> const&, Precision)
{
  return kernel_types_yes;
}
 
template <MutableRankedMdspecLike<2> A, MutableRankedMdspecLike<2> B>
  requires HostWritableMdspec<A> && HostWritableMdspec<B> && RealOrComplex<typename A::value_type> &&
           has_runtime_precision_v<typename A::value_type> &&
           std::same_as<typename A::value_type, typename B::value_type>
KernelAttempt try_kernel(CpuReferenceBackend, linear_solve_op const&, A& a, B& b, SolveInfo& info,
                         SolveOptions<make_real_t<typename A::value_type>> const& options, Precision p)
{
  using R = make_real_t<typename A::value_type>;
  (void)p.bit_count();
  ERROR_IF(a.extent(0) != a.extent(1) || a.extent(0) != b.extent(0), "solve operand shapes do not agree");
  detail::require_solve_options(options);
  info = {};
  if (a.extent(0) == 0 || b.extent(1) == 0) return KernelAttempt::success;
  auto aa = acquire_host_write_access_sync(a);
  auto ba = acquire_host_write_access_sync(b);
  auto as = aa.mdspan();
  auto bs = ba.mdspan();
  if (!detail::solve_matrix_is_finite(as) || !detail::solve_matrix_is_finite(bs))
  {
    info.status = SolveStatus::nonfinite_input;
    return KernelAttempt::success;
  }
  auto av = packed_detail::pack(as, p), bv = packed_detail::pack(bs, p);
  auto aw = av.mdspan();
  auto bw = bv.mdspan();
  if (!detail::solve_matrix_is_finite(aw) || !detail::solve_matrix_is_finite(bw))
  {
    info.status = SolveStatus::nonfinite_result;
    return KernelAttempt::success;
  }
  SolveOptions<R> converted{.relative_pivot_tolerance = scalar_at_precision<R>(options.relative_pivot_tolerance, p)};
  std::vector<std::size_t> pivots(std::size_t(a.extent(0)));
  detail::cpu_reference::lu_factor(aw, std::span(pivots), nullptr, info, converted, scalar_at_precision<R>(0, p));
  packed_detail::unpack(av, as);
  if (info.succeeded())
  {
    detail::cpu_reference::lu_solve(aw, std::span<std::size_t const>(pivots), bw, info);
    packed_detail::unpack(bv, bs);
  }
  return KernelAttempt::success;
}
 
} // namespace uni20::linalg