4  Implementation of the Mixed IR Algorithm

4.1 Review of LU Decomposition

Partial Pivoting

In partial pivoting only rows are possibly permuted and the permutations are ‘stored’ in a matrix \(P\). So we compute the decomposition:

\[ LU = PA \]

This can be used to solve the matrix equation \(Ax = b\) as follows:

\[ \begin{aligned} &\quad Ax = b \\ &\Leftrightarrow (PA)x = Pb \\ &\Leftrightarrow L(Ux) = Pb &&\text{($LU = PA$, assoc)} \\ &\Rightarrow Ly = Pb &&\text{($y := Ux$)} \\ &\text{solve for $y$ using forward subst} \\ &\text{solve $Ux = y$ for $x$ using backward subst} \end{aligned} \]

Full Pivoting

In full pivoting columns can be permuted in addition to the row permutation, improving backward stability of the algorithm, s.t. the following equation holds:

\[ PAQ = LU \]

Here \(P\), \(Q\) are the row and column permutation matrices. This decomposition can be used to solve the equation \(Ax = b\) as follows:

\[ \begin{aligned} &\quad Ax = b \\ &\Leftrightarrow PA (QQ^{-1})x = b \\ &\Leftrightarrow (PAQ)Q^{-1}x = b \\ &\Leftrightarrow L(UQ^{-1})x = b &&\text{($LU = PAQ$)} \\ &\Rightarrow Lz = b &&\text{($UQ^{-1} := z$)}\\ &\text{Solve for $z$ using forward subst.}\\ &\text{We have:} \\ &\Rightarrow Uy = z &&\text{($Q^{-1}x := y$)} \\ &\text{Solve for $y$ using backward subst.} \\ &\text{Compute $x = Qy$} \end{aligned} \]

Implementation in HDNum

Solving \(Ax = b\) using \(LU\) decomposition is implemented in HDnum as follows. Consider:

\[ A = \begin{bmatrix} 2 & 1 & 7 \\ 8 & 8 & 33 \\ -4 & 10 & 4 \end{bmatrix}, \qquad b = \begin{bmatrix} 15 \\ 73 \\ 12 \end{bmatrix} \]

We enter this as follows:

hdnum::Vector<hdnum::FP256> b(3);
b[0] = 15;
b[1] = 73;
b[2] = 12;

hdnum::DenseMatrix <hdnum::FP256> A(3, 3);
A[0][0] = 2;    A[0][1] = 1;    A[0][2] = 7;
A[1][0] = 8;    A[1][1] = 8;    A[1][2] = 33;
A[2][0] = -4;   A[2][1] = 10;   A[2][2] = 4;

We also need the vector \(x\) to hold the solution, and the vectors \(p\) and \(q\) to hold the row and columnt permutations:

hdnum::Vector<hdnum::FP256> x(3, 0.0); // holds the solution
hdnum::Vector<std::size_t> p(3);
hdnum::Vector<std::size_t> q(3);

Then we carry out the \(LU\) decomposition as follows:

hdnum::lr_fullpivot(A, p, q);

After this command \(A\) is modified and holds both \(L\) and \(U\), while \(p\) and \(q\) hold the permutation information. We also have to permute \(b\) accordingly:

hdnum::permute_forward(p, b);

Then we solve the \(Ly = b\) part of the equation (the forward substitution):

hdnum::solveL(A, b, b);

Note that we are ovewriting \(b\) and storing the result directly in it to save space. Then, we finally complete the solution with:

hdnum::solveR(A, x, b);

Since we did full-pivoting, we should permute back the solution with:

hdnum::permute_backward(q, x);

Now \(x\) holds the solution.

4.2 The Mixed IR Algorithm

First we remember the mixed-ir algorithm:

1. Compute LU factorization PAQ = LU in precision uf.
2. Solve LU x0 = P b in precision uf;
   Cast x0 to precision u
3. While not converged:
   4. Compute ri = b - A xi in precision ur (residual)
   5. cast ri to precision u.
   6. Solve LU di = P ri in precision uf;
      Cast di to precision u
   7. update x(i+1)= xi + di in precision u
  • Convergence condition \(\frac{\norm{d_i}}{\norm{x_i}} < u\), where \(u\) is the unit roundoff of the working precision, or stop after a max number of iterations.

The c++ implementation of this algorithm is given as follows. First we define an options struct that holds the options that control max iterations and relative correction tolerance, i.e. the termination condition.

struct MixedIROptions {
    std::size_t max_iterations = 10;
    double rel_correction_tol = 0.0;
};

The default value rel_correctin_tol = 0.0 causes the algorithm to try to use the default unit roundoff of the data type used for the working precision \(u\), as we’ll see later in the algorithm.

Next we define a struct to hold the resulst, here the comments provide the explanations:

template<class T_work>
struct MixedIRResult {
    // Final approximate solution, stored in the working precision T_work.
    hdnum::Vector<T_work> x;

    // Number of refinement iterations actually performed.
    // This does not count the initial LU solve for x0.
    std::size_t iterations = 0;

    // True if the stopping criterion was reached before max_iterations.
    // False means the algorithm stopped because the iteration limit was reached.
    bool converged = false;

    // Last computed relative correction norm ||d_i|| / ||x_i||.
    // Useful for diagnostics, especially if converged == false.
    double final_rel_correction = 0.0;

    // History of relative correction norms, one entry per refinement iteration.
    // Useful for convergence diagnostics and later plotting.
    std::vector<double> rel_corrections;
};

We define a function that returns \(u\) of the templatized type T, if it’s obtainable, or throws an error if not. (Mostly we will use double for \(u\), for which it can be computed in compile time):

template<class T>
double default_unit_roundoff()
{
    static_assert(std::numeric_limits<T>::is_specialized,
        "default_unit_roundoff: no numeric_limits specialization for this type;"
         "please set_rel_correction_tol_explicitly");
    
    return 0.5 * static_cast<double>(std::numeric_limits<T>::epsilon());
}

Next a function that computes the 2-norm in double precision for a vector of a generic T type.

template<class T>
double norm2_as_double(const hdnum::Vector<T>& v)
{
    double sum = 0.0;

    for (std::size_t i = 0; i < v.size(); ++i) {
        const double vi = scalar_cast<double>(v[i]);
        sum += vi * vi;
    }

    return std::sqrt(sum);
}

Then we define a function that encapsulates the steps to solve an equation \(Ax = b\) after \(LU\) decomposition is done. As input it expects the result of the \(LU\) decomposition, the permutation vectors \(p\) and \(q\), and the vector \(b\):

template<class T>
hdnum::Vector<T>
solve_with_lu_fullpivot(
    const hdnum::DenseMatrix<T>& LU,
    const hdnum::Vector<std::size_t>& p,
    const hdnum::Vector<std::size_t>& q,
    const hdnum::Vector<T>& b
)
{
    hdnum::Vector<T> y(b); // copy b 
    hdnum::Vector<T> x(b.size()); // holds solution

    hdnum::permute_forward(p, y);
    hdnum::solveL(LU, y, y); // Solve L*y = P*b, in-place in y.
    hdnum::solveR(LU, x, y); // Solve U*x = y
    hdnum::permute_backward(q, x); // undo column permutation
    return x;
}

Then a function that performs the residual computation with the necessary type conversions. It expects a matrix \(A\), vectors \(b\) and current solution \(x\), all in working precision \(u\), converts them to higher precisions, and returns the residual \(r = Ax - b\). in higher residual precision \(u_r\):

template<class T_residual, class T_work>
hdnum::Vector<T_residual> compute_residual(
    const hdnum::DenseMatrix<T_work>& A,
    const hdnum::Vector<T_work>& b,
    const hdnum::Vector<T_work>& x
)
{
    const std::size_t n = b.size();

    hdnum::DenseMatrix<T_residual> A_r(A.rowsize(), A.colsize());
    hdnum::Vector<T_residual> b_r(n);
    hdnum::Vector<T_residual> x_r(n);
    hdnum::Vector<T_residual> Ax_r(n);
    hdnum::Vector<T_residual> r_r(n);

    convert(A_r, A);
    convert(b_r, b);
    convert(x_r, x);

    A_r.mv(Ax_r, x_r);

    for (std::size_t i = 0; i < n; i++) {
        r_r[i] = b_r[i] - Ax_r[i];
    }

    return r_r;
}

Finally we provide the full algorithm, where the comments provide the necessary explanations:

template <class T_factor, class T_work, class T_residual>
MixedIRResult<T_work> mixed_ir(
    const hdnum::DenseMatrix<T_work>& A,
    const hdnum::Vector<T_work>& b,
    const MixedIROptions& options = {}
)
{
    if (A.rowsize() != A.colsize() || A.rowsize() == 0) {
        throw std::invalid_argument("mixed_ir: A must be square and non-empty");
    }

    if (A.rowsize() != b.size()) {
        throw std::invalid_argument("mixed_ir: A must be square and nonempty");
    }

    const std::size_t n = b.size();

    // if tol is default value 0.0, then use the unit roundoff of T_work
    // which is compute in compile time
    // otherwise use the value passed in the options. 
    const double tol = 
        options.rel_correction_tol > 0.0 ? options.rel_correction_tol
                                         : default_unit_roundoff<T_work>();

    // 1. convert A and b to factorization precision.
    hdnum::DenseMatrix<T_factor> A_f(n, n);
    hdnum::Vector<T_factor> b_f(n);

    convert(A_f, A);
    convert(b_f, b);

    // 2. Compute LU factorization in low precision
    //
    // lr_fullpivot overwrites A_f with the combined L / U factors
    hdnum::Vector<std::size_t> p(n);
    hdnum::Vector<std::size_t> q(n);

    hdnum::lr_fullpivot(A_f, p, q);


    // 3. Initial low-precision solve
    hdnum::Vector<T_factor> x0_f = solve_with_lu_fullpivot(A_f, p, q, b_f);

    // 4. cast x0 to working precision
    MixedIRResult<T_work> result;
    result.x = hdnum::Vector<T_work>(n);
    convert(result.x, x0_f);

    // 5. refinement loop
    for (std::size_t k = 0; k < options.max_iterations; k++) {
        // Compute residual r = b - A * x in residual precision.
        hdnum::Vector<T_residual> r_r = compute_residual<T_residual>(A, b, result.x);

        // cast residual to working precision
        hdnum::Vector<T_work> r_w(n);
        convert(r_w, r_r);

        // The triangular solve uses T_factor L and U factors,
        // so the right-hand side must also be in T_factor
        hdnum::Vector<T_factor> r_f(n);
        convert(r_f, r_w);

        // Solve correction equation A d = r using existing LU factors:
        hdnum::Vector<T_factor> d_f = solve_with_lu_fullpivot(A_f, p, q, r_f);

        // Cast correction back to working precision
        hdnum::Vector<T_work> d_w(n);
        convert(d_w, d_f);

        const double norm_d = norm2_as_double(d_w);
        const double norm_x = norm2_as_double(result.x);
        const double denominator = (norm_x > 0.0) ? norm_x : 1.0;

        const double rel_corr = norm_d / denominator;

        result.rel_corrections.push_back(rel_corr);
        result.final_rel_correction = rel_corr;

        // update in working precision
        for (std::size_t i = 0; i < n; i++) {
            result.x[i] += d_w[i];
        }

        result.iterations = k + 1;

        if (norm_d <= tol * norm_x) {
            result.converged = true;
            break;
        }

    }

    return result;
}

4.3 Unit Roundoff \(u\)

For a floating point format with \(t\) significand bits the unit round off is defined as

\[ \begin{aligned} u &= \frac{1}{2}\beta^{1-t} \\ &= 2^{-t} &&(\beta = 2) \end{aligned} \]

hdnum’s cpfloat wrapper’s are:

using FP16 = hdnum::CPFloat<11,5>;
using bfloat16 = hdnum::CPFloat<8,8>;
using FP8 = hdnum::CPFloat<4,4>;

Where the first template parameters are the significand (or the mantisse) bits.

and the FP wrappers:

using FP128  = FP<64>;
using FP256  = FP<192>;
using FP512  = FP<448>;
using FP1024 = FP<960>;

where the template parameters are the mantisse bits.

From here we can construct the following table summarizing the unit roundoff \(u\) for the various precisions:

HDNUM type Actual HDNUM alias Significand precision (t) Unit roundoff \(u = 2^{-t}\)
FP8 CPFloat<4,4> 4 \(2^{-4} \approx 6.25 \cdot 10^{-2}\)
bfloat16 CPFloat<8,8> 8 \(2^{-8} \approx 3.91 \cdot 10^{-3}\)
FP16 CPFloat<11,5> 11 \(2^{-11} \approx 4.88 \cdot 10^{-4}\)
FP32 float 24 \(2^{-24} \approx 5.96 \cdot 10^{-8}\)
FP64 double 53 \(2^{-53} \approx 1.11 \cdot 10^{-16}\)
FP128 FP<64> 64 \(2^{-64} \approx 5.42 \cdot 10^{-20}\)
FP256 FP<192> 192 \(2^{-192} \approx 1.59 \cdot 10^{-58}\)
FP512 FP<448> 448 \(2^{-448} \approx 1.38 \cdot 10^{-135}\)
FP1024 FP<960> 960 \(2^{-960} \approx 1.03 \cdot 10^{-289}\)

unit roundoff traits

Orignially we had the function

template<class T>
double default_unit_roundoff()
{
    static_assert(std::numeric_limits<T>::is_specialized,
        "default_unit_roundoff: no numeric_limits specialization for this type;"
         "please set_rel_correction_tol_explicitly");
    
    return 0.5 * static_cast<double>(std::numeric_limits<T>::epsilon());
}

This works only for data types float and double. We improve this using type traits, s.t. unit roundoffs are returned for other floating point formats at compile time:

#pragma once

#include <limits>
#include <type_traits>

#include "hdnum.hh"

namespace mpir {

constexpr double pow2_neg(int p)
{
    double x = 1.0;
    for (int i = 0; i < p; ++i) {
        x *= 0.5;
    }
    return x;
}

namespace detail {

template<class>
struct unit_roundoff_always_false : std::false_type {};

} // namespace detail

template<class T>
struct unit_roundoff_traits {
    static_assert(detail::unit_roundoff_always_false<T>::value,
        "unit_roundoff_traits: unsupported scalar type");
};

template<>
struct unit_roundoff_traits<hdnum::FP32> {
    static constexpr double value =
        0.5 * static_cast<double>(std::numeric_limits<hdnum::FP32>::epsilon());
};

template<>
struct unit_roundoff_traits<hdnum::FP64> {
    static constexpr double value =
        0.5 * static_cast<double>(std::numeric_limits<hdnum::FP64>::epsilon());
};

#ifdef HDNUM_HAS_CPFLOAT
template<int m, int e>
struct unit_roundoff_traits<hdnum::CPFloat<m,e>> {
    static constexpr double value = pow2_neg(m);
};
#endif

#ifdef HDNUM_HAS_GMP
template<int m>
struct unit_roundoff_traits<hdnum::FP<m>> {
    static constexpr double value = pow2_neg(m);
};
#endif

template<class T>
constexpr double default_unit_roundoff()
{
    return unit_roundoff_traits<T>::value;
}

} // namespace mpir

Now mixed ir algorithm will work for all fp formats without having to specify tol manually in the options struct.