AMReX_Smoother_MV.H

Go to the documentation of this file.

#ifndef AMREX_SMOOTHER_MV_H_
#define AMREX_SMOOTHER_MV_H_
 
#include <AMReX_Algebra.H>
#include <AMReX_Array.H>
#include <AMReX_BLProfiler.H>
#include <AMReX_OpenMP.H>
#include <cmath>
#include <memory>
#include <utility>
 
namespace amrex {
 
template <typename T>
class JacobiSmoother
{
public:
    explicit JacobiSmoother (SpMatrix<T> const* a_A, bool a_l1 = false)
        : m_A(a_A), m_l1(a_l1), m_weight(a_l1 ? T(1) : T(2./3.)) {}
 
    int setNumIters (int a_niters) { return std::exchange(m_niters, a_niters); }
 
    void setWeight (T w) { m_weight = w; }
 
    void setL1 (bool b) { m_l1 = b; m_weight = b ? T(1) : T(2./3.); m_scaling.reset(); }
 
    void operator() (AlgVector<T>& xvec, AlgVector<T> const& bvec,
                     bool with_initial_guess = false)
    {
        BL_PROFILE("JacobiSmoother");
        AMREX_ASSERT(xvec.partition() == m_A->partition());
        auto const& diag = scaling();
        if ( ! with_initial_guess) {
            xvec.setValAsync(0);
        }
        T const w = m_weight;
        for (int iter = 0; iter < m_niters; ++iter) {
            if ((iter == 0) && ! with_initial_guess) {
                // A x is zero, so this sweep is just x = w b / d.
                ForEach(xvec, bvec, diag,
                        [=] AMREX_GPU_DEVICE (T& x, T const& b, T const& d)
                        {
                            if (d != T(0)) {
                                x += w * b/d;
                            }
                        });
            } else {
                auto& Axvec = workspace(xvec.partition()); // allocated on first need
                SpMV(Axvec, *m_A, xvec);
                ForEach(xvec, Axvec, bvec, diag,
                        [=] AMREX_GPU_DEVICE (T& x, T const& ax, T const& b, T const& d)
                        {
                            if (d != T(0)) {
                                x += w * (b-ax)/d;
                            }
                        });
            }
        }
        Gpu::streamSynchronize();
    }
 
private:
    SpMatrix<T> const* m_A;
    bool m_l1 = false;
    T m_weight;
    int m_niters = 4;
    // l1 row norms, built on first use. Shared so that the smoother stays
    // copyable (it is used as a std::function preconditioner).
    std::shared_ptr<AlgVector<T>> m_scaling;
    std::shared_ptr<AlgVector<T>> m_work;
 
    AlgVector<T>& workspace (AlgPartition const& partition)
    {
        if (!m_work || m_work->partition() != partition) {
            m_work = std::make_shared<AlgVector<T>>(partition);
        }
        return *m_work;
    }
 
public: // NOLINT  Private, but public for CUDA
    AlgVector<T> const& scaling ()
    {
        if (!m_l1) { return m_A->diagonalVector(); }
        if (!m_scaling) {
            // The matrix might not have been split yet (e.g., no SpMV so far).
            const_cast<SpMatrix<T>*>(m_A)->setColumnPartition(m_A->partition());
            m_scaling = std::make_shared<AlgVector<T>>(m_A->partition());
            auto* AMREX_RESTRICT p = m_scaling->data();
            auto const& a = m_A->const_parcsr();
            auto const& d = m_A->diagonalVector();
            auto const* AMREX_RESTRICT pd = d.data();
            ParallelForOMP(m_scaling->numLocalRows(), [=] AMREX_GPU_DEVICE (Long i) noexcept
            {
                T s = T(0);
                for (auto idx = a.csr0.row_offset[i]; idx < a.csr0.row_offset[i+1]; ++idx) {
                    s += std::abs(a.csr0.mat[idx]);
                }
                if (detail::has_remote_row(a, i)) {
                    auto ii = a.row_map[i];
                    for (auto idx = a.csr1.row_offset[ii]; idx < a.csr1.row_offset[ii+1]; ++idx) {
                        s += std::abs(a.csr1.mat[idx]);
                    }
                }
                p[i] = std::copysign(s, pd[i]);
            });
            Gpu::streamSynchronize();
        }
        return *m_scaling;
    }
};
 
template <typename T>
class ChebyshevSmoother
{
public:
    explicit ChebyshevSmoother (SpMatrix<T> const* a_A, bool a_l1 = false)
        : m_jacobi(a_A, a_l1), m_A(a_A), m_l1(a_l1) {}
 
    void setDegree (int d) { m_degree = d; }
    void setEigRatio (T r) {
        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(r > T(1), "ChebyshevSmoother: the ratio must exceed 1");
        m_ratio = r;
    }
    int setNumIters (int a_niters) { return std::exchange(m_niters, a_niters); }
    void setLambdaMax (T l) { m_lambda_max = l; }
    T lambdaMax ()
    {
        if (m_lambda_max <= T(0)) { bound(); }
        return m_lambda_max;
    }
 
    void operator() (AlgVector<T>& xvec, AlgVector<T> const& bvec,
                     bool with_initial_guess = false)
    {
        BL_PROFILE("ChebyshevSmoother");
        AMREX_ASSERT(xvec.partition() == m_A->partition());
        T const lmax = lambdaMax();
        T const lmin = lmax / m_ratio;
        T const theta = (lmax + lmin) / T(2);
        T const delta = (lmax - lmin) / T(2);
        T const sigma = theta / delta;
 
        auto const& diag = m_jacobi.scaling();
        if (!m_work || (*m_work)[0].partition() != xvec.partition()) {
            m_work = std::make_shared<Array<AlgVector<T>,3>>();
            for (auto& v : *m_work) { v.define(xvec.partition()); }
        }
        auto& r  = (*m_work)[0];
        auto& d  = (*m_work)[1];
        auto& ad = (*m_work)[2];
        if ( ! with_initial_guess) { xvec.setValAsync(0); }
 
        T const c0 = T(1) / theta;
        for (int it = 0; it < m_niters; ++it) {
            // r = b - A x; d = D^{-1} r / theta; x += d
            if (it == 0 && ! with_initial_guess) {
                ForEach(xvec, d, r, bvec, diag,
                        [=] AMREX_GPU_DEVICE (T& xi, T& di, T& ri, T const& bi,
                                              T const& dd)
                        {
                            ri = bi;
                            di = (dd != T(0)) ? c0 * ri / dd : T(0);
                            xi += di;
                        });
            } else {
                SpMV(r, *m_A, xvec);
                ForEach(xvec, d, r, bvec, diag,
                        [=] AMREX_GPU_DEVICE (T& xi, T& di, T& ri, T const& bi,
                                              T const& dd)
                        {
                            ri = bi - ri;
                            di = (dd != T(0)) ? c0 * ri / dd : T(0);
                            xi += di;
                        });
            }
            T rho_old = T(1) / sigma;
            for (int k = 1; k < m_degree; ++k) {
                // r -= A d; rho = 1/(2 sigma - rho_old);
                // d = rho rho_old d + (2 rho / delta) D^{-1} r; x += d
                SpMV(ad, *m_A, d);
                T const rho = T(1) / (T(2)*sigma - rho_old);
                T const c1 = rho * rho_old;
                T const c2 = T(2) * rho / delta;
                ForEach(xvec, d, r, ad, diag,
                        [=] AMREX_GPU_DEVICE (T& xi, T& di, T& ri, T const& adi,
                                              T const& dd)
                        {
                            ri -= adi;
                            di = c1 * di + ((dd != T(0)) ? c2 * ri / dd : T(0));
                            xi += di;
                        });
                rho_old = rho;
            }
        }
        Gpu::streamSynchronize();
    }
 
private:
    JacobiSmoother<T> m_jacobi; // supplies the scaling vector
    SpMatrix<T> const* m_A;
    bool m_l1 = false;
    int m_degree = 2;
    T m_ratio = T(20);
    int m_niters = 1;
    T m_lambda_max = T(-1);
    std::shared_ptr<Array<AlgVector<T>,3>> m_work; // r, d, A d
 
public: // NOLINT  Private, but public for CUDA
    void bound ()
    {
        if (m_l1) { m_lambda_max = T(1); return; }
        // The matrix might not have been split yet (e.g., no SpMV so far).
        const_cast<SpMatrix<T>*>(m_A)->setColumnPartition(m_A->partition());
        auto const& a = m_A->const_parcsr();
        auto const* AMREX_RESTRICT pd = m_A->diagonalVector().data();
        T lmax = Reduce::Max<T>(m_A->numLocalRows(), [=] AMREX_GPU_DEVICE (Long i) -> T
        {
            if (pd[i] == T(0)) { return T(0); }
            T s = T(0); // sum_j |a_ij|
            for (auto idx = a.csr0.row_offset[i]; idx < a.csr0.row_offset[i+1]; ++idx) {
                s += std::abs(a.csr0.mat[idx]);
            }
            if (detail::has_remote_row(a, i)) {
                auto ii = a.row_map[i];
                for (auto idx = a.csr1.row_offset[ii]; idx < a.csr1.row_offset[ii+1]; ++idx) {
                    s += std::abs(a.csr1.mat[idx]);
                }
            }
            return s / std::abs(pd[i]);
        });
        ParallelAllReduce::Max(lmax, ParallelContext::CommunicatorSub());
        m_lambda_max = (lmax > T(0)) ? lmax : T(1);
    }
};
 
template <typename T>
class L1GaussSeidelSmoother
{
public:
    explicit L1GaussSeidelSmoother (SpMatrix<T> const* a_A) : m_A(a_A) {}
 
    int setNumIters (int a_niters) { return std::exchange(m_niters, a_niters); }
 
    void operator() (AlgVector<T>& xvec, AlgVector<T> const& bvec,
                     bool with_initial_guess = false, bool backward = false)
    {
#ifdef AMREX_USE_GPU
        amrex::ignore_unused(xvec, bvec, with_initial_guess, backward);
        amrex::Abort("L1GaussSeidelSmoother is only available in CPU builds");
#else
        BL_PROFILE("L1GaussSeidelSmoother");
        AMREX_ASSERT(xvec.partition() == m_A->partition());
        auto& A = *const_cast<SpMatrix<T>*>(m_A);
        // The matrix might not have been split yet (e.g., no SpMV so far).
        A.setColumnPartition(A.partition());
        auto const& a = A.const_parcsr();
        Long const nrows = xvec.numLocalRows();
        T const* AMREX_RESTRICT dl1 = l1_diagonal().data();
        int const nblocks = m_nblocks;
        bool const has_remote = this->has_remote();
        T const* AMREX_RESTRICT b = bvec.data();
        T* AMREX_RESTRICT x = xvec.data();
 
        // The previous iterate is needed for couplings between row blocks,
        // and the remote part of A x for couplings to other processes.
        if (!m_work) {
            m_work = std::make_shared<Array<AlgVector<T>,2>>();
            if (nblocks > 1) { (*m_work)[0].define(xvec.partition()); }
            if (has_remote) { (*m_work)[1].define(xvec.partition()); }
        }
        auto& xold = (*m_work)[0];
        auto& yrem = (*m_work)[1];
        T const* AMREX_RESTRICT px_old = (nblocks > 1) ? xold.data() : nullptr;
        T const* AMREX_RESTRICT py_rem = has_remote ? yrem.data() : nullptr;
 
        if ( ! with_initial_guess) { xvec.setVal(0); }
 
        // x_i += (b_i - (A x)_i) / d_i, row by row, with d_i = a_ii + l1 term.
        auto sweep = [&] (auto multi_block, auto remote)
        {
#ifdef AMREX_USE_OMP
#pragma omp parallel for
#endif
            for (int iblock = 0; iblock < nblocks; ++iblock) {
                auto const [rlo, rhi] = row_block(nrows, iblock, nblocks);
                auto const& c0 = a.csr0;
                for (Long n = 0; n < rhi-rlo; ++n) {
                    Long const i = backward ? (rhi-1-n) : (rlo+n);
                    T r = b[i];
                    if constexpr (remote) { r -= py_rem[i]; }
                    for (Long idx = c0.row_offset[i]; idx < c0.row_offset[i+1]; ++idx) {
                        Long const j = c0.col_index[idx];
                        if constexpr (multi_block) {
                            // Rows of this block use the current iterate, others the old one.
                            r -= c0.mat[idx] * ((j >= rlo && j < rhi) ? x[j] : px_old[j]);
                        } else {
                            r -= c0.mat[idx] * x[j];
                        }
                    }
                    if (dl1[i] != T(0)) { x[i] += r / dl1[i]; }
                }
            }
        };
 
        for (int iter = 0; iter < m_niters; ++iter) {
            bool const zero_guess = (iter == 0 && ! with_initial_guess);
            if (nblocks > 1) {
                if (zero_guess) { xold.setVal(0); } else { xold.copy(xvec); }
            }
            if (has_remote) {
                // Remote contribution with the previous iterate.
                yrem.setVal(0);
                if (!zero_guess) {
                    A.startComm_mv(xvec);
                    A.finishComm_mv(yrem);
                }
            }
            if (nblocks > 1) {
                if (has_remote) { sweep(std::true_type{}, std::true_type{}); }
                else            { sweep(std::true_type{}, std::false_type{}); }
            } else {
                if (has_remote) { sweep(std::false_type{}, std::true_type{}); }
                else            { sweep(std::false_type{}, std::false_type{}); }
            }
        }
#endif
    }
 
private:
    SpMatrix<T> const* m_A;
    int m_niters = 2;
    int m_nblocks = -1;
    int m_has_remote = -1; // any rank couples to another rank; -1: not yet known
    std::shared_ptr<AlgVector<T>> m_dl1; // a_ii + l1 term for the current row blocks
    std::shared_ptr<Array<AlgVector<T>,2>> m_work; // previous iterate, remote part
 
    bool has_remote ()
    {
        if (m_has_remote < 0) {
            int r = (m_A->const_parcsr().csr1.nnz > 0) ? 1 : 0;
            ParallelAllReduce::Max(r, ParallelContext::CommunicatorSub());
            m_has_remote = r;
        }
        return m_has_remote == 1;
    }
 
    static std::pair<Long,Long> row_block (Long nrows, int iblock, int nblocks)
    {
        Long const lo = (nrows * iblock) / nblocks;
        Long const hi = (nrows * (iblock+1)) / nblocks;
        return {lo, hi};
    }
 
    template <typename C>
    static T diag_of (C const& c0, Long i)
    {
        for (Long idx = c0.row_offset[i]; idx < c0.row_offset[i+1]; ++idx) {
            if (c0.col_index[idx] == i) { return c0.mat[idx]; }
        }
        return T(0);
    }
 
#ifndef AMREX_USE_GPU
    AlgVector<T> const& l1_diagonal ()
    {
        int const nb = OpenMP::get_max_threads();
        if (!m_dl1 || nb != m_nblocks) {
            m_nblocks = nb;
            m_work.reset(); // its vectors depend on the number of blocks
            m_dl1 = std::make_shared<AlgVector<T>>(m_A->partition());
            T* AMREX_RESTRICT p = m_dl1->data();
            auto const& a = m_A->const_parcsr();
            Long const nrows = m_dl1->numLocalRows();
#ifdef AMREX_USE_OMP
#pragma omp parallel for
#endif
            for (int iblock = 0; iblock < nb; ++iblock) {
                auto const [rlo, rhi] = row_block(nrows, iblock, nb);
                auto const& c0 = a.csr0;
                for (Long i = rlo; i < rhi; ++i) {
                    T const aii = diag_of(c0, i); // same entry the sweep uses
                    T l1 = T(0);
                    for (Long idx = c0.row_offset[i]; idx < c0.row_offset[i+1]; ++idx) {
                        Long const j = c0.col_index[idx];
                        if (j != i && (j < rlo || j >= rhi)) {
                            l1 += std::abs(c0.mat[idx]);
                        }
                    }
                    if (detail::has_remote_row(a, i)) {
                        Long const ii = a.row_map[i];
                        for (Long idx = a.csr1.row_offset[ii]; idx < a.csr1.row_offset[ii+1]; ++idx) {
                            l1 += std::abs(a.csr1.mat[idx]);
                        }
                    }
                    p[i] = aii + std::copysign(l1, aii);
                }
            }
        }
        return *m_dl1;
    }
#endif
};
 
}
 
#endif