AMReX_CSR.H

Go to the documentation of this file.

#ifndef AMREX_CSR_H_
#define AMREX_CSR_H_
#include <AMReX_Config.H>
 
#include <AMReX_Gpu.H>
#include <AMReX_INT.H>
#include <AMReX_OpenMP.H>
 
#if defined(AMREX_USE_CUDA)
#include <cub/cub.cuh> // for Clang
#endif
 
#include <algorithm>
#include <concepts>
#include <climits>
#include <type_traits>
 
namespace amrex {
 
template <typename T, typename I = Long>
struct CsrView {
    using index_type = I;
    using CI = std::conditional_t<std::is_const_v<T>, I const, I>;
    T* AMREX_RESTRICT mat = nullptr;
    CI* AMREX_RESTRICT col_index = nullptr;
    CI* AMREX_RESTRICT row_offset = nullptr;
    Long nnz = 0;
    Long nrows = 0;
};
 
template <typename T, template <typename> class V, typename I = Long>
struct CSR {
    using index_type = I;
    V<T> mat;
    V<I> col_index;
    V<I> row_offset;
    Long nnz = 0;
 
    [[nodiscard]] Long nrows () const {
        return row_offset.empty() ? Long(0) : Long(row_offset.size())-1;
    }
 
    void resize (Long num_rows, Long num_non_zeros) {
        mat.resize(num_non_zeros);
        col_index.resize(num_non_zeros);
        row_offset.resize(num_rows+1);
        nnz = num_non_zeros;
    }
 
    CsrView<T,I> view () {
        return CsrView<T,I>{mat.data(), col_index.data(), row_offset.data(),
                            nnz, Long(row_offset.size())-1};
    }
 
    [[nodiscard]] CsrView<T const,I> view () const {
        return CsrView<T const,I>{.mat = mat.data(),
                                  .col_index = col_index.data(),
                                  .row_offset = row_offset.data(),
                                  .nnz = nnz,
                                  .nrows = Long(row_offset.size())-1};
    }
 
    [[nodiscard]] CsrView<T const,I> const_view () const {
        return CsrView<T const,I>{mat.data(), col_index.data(), row_offset.data(),
                                  nnz, Long(row_offset.size())-1};
    }
 
    void sort ();
 
    void sort_on_host ();
};
 
template <typename C, typename T, template<typename> class AD, template<typename> class AS,
          typename I>
requires (std::same_as<C,Gpu::HostToDevice> ||
          std::same_as<C,Gpu::DeviceToHost> ||
          std::same_as<C,Gpu::DeviceToDevice>)
void duplicateCSR (C c, CSR<T,AD,I>& dst, CSR<T,AS,I> const& src)
{
    dst.mat.resize(src.mat.size());
    dst.col_index.resize(src.col_index.size());
    dst.row_offset.resize(src.row_offset.size());
    Gpu::copyAsync(c,
                   src.mat.begin(),
                   src.mat.end(),
                   dst.mat.begin());
    Gpu::copyAsync(c,
                   src.col_index.begin(),
                   src.col_index.end(),
                   dst.col_index.begin());
    Gpu::copyAsync(c,
                   src.row_offset.begin(),
                   src.row_offset.end(),
                   dst.row_offset.begin());
    dst.nnz = src.nnz;
}
 
template <typename T, template <typename> class V, typename I>
void CSR<T,V,I>::sort ()
{
    if (nnz <= 0) { return; }
 
#ifdef AMREX_USE_GPU
 
#if defined(AMREX_USE_CUDA) || defined(AMREX_USE_HIP)
 
    // The function is synchronous. If that is no longer the case, we might
    // need to update SpMatrix::define.
 
    constexpr int nthreads = 256;
    constexpr int nwarps_per_block = nthreads / Gpu::Device::warp_size;
 
    AMREX_ALWAYS_ASSERT((nrows()+nwarps_per_block-1) < Long(std::numeric_limits<int>::max()));
 
    auto nr = int(nrows());
    int nblocks = (nr + nwarps_per_block-1) / nwarps_per_block;
    auto const& stream = Gpu::gpuStream();
 
    auto* pmat = mat.data();
    auto* pcol = col_index.data();
    auto* prow = row_offset.data();
 
    // Short rows (the common case on fine AMG levels): one thread per row,
    // insertion sort in registers. Longer rows are left to the warp kernel
    // below. Both kernels skip rows that are already sorted.
    // The cap keeps the per-thread arrays at ~256 bytes for wide value types.
    constexpr int SHORT_ROW = std::clamp(int(256/(sizeof(I)+sizeof(T))), 4, 16);
    ParallelFor(nr, [=] AMREX_GPU_DEVICE (int r) noexcept
    {
        Long const b = prow[r];
        int const len = int(prow[r+1] - b);
        if (len <= 1 || len > SHORT_ROW) { return; }
        I keys[SHORT_ROW];
        bool sorted = true;
        for (int i = 0; i < len; ++i) {
            keys[i] = pcol[b+i];
            if (i > 0) { sorted = sorted && (keys[i-1] <= keys[i]); }
        }
        if (sorted) { return; }
        T vals[SHORT_ROW];
        for (int i = 0; i < len; ++i) {
            vals[i] = pmat[b+i];
        }
        for (int i = 1; i < len; ++i) {
            I const k = keys[i];
            T const v = vals[i];
            int j = i - 1;
            while (j >= 0 && keys[j] > k) {
                keys[j+1] = keys[j];
                vals[j+1] = vals[j];
                --j;
            }
            keys[j+1] = k;
            vals[j+1] = v;
        }
        for (int i = 0; i < len; ++i) {
            pcol[b+i] = keys[i];
            pmat[b+i] = vals[i];
        }
    });
 
    Gpu::Buffer<int> needs_fallback({0});
    auto* d_needs_fallback = needs_fallback.data();
 
    amrex::launch_global<nthreads><<<nblocks, nthreads, 0, stream>>>
        ([=] AMREX_GPU_DEVICE () noexcept
    {
        int wid = int(threadIdx.x)/Gpu::Device::warp_size;
        int r = int(blockIdx.x)*nwarps_per_block + wid;
        if (r >= nr) return;
 
        Long const b = prow[r];
        Long const e = prow[r+1];
        auto const len = int(e - b);
 
        if (len <= 1) return;
 
        int lane = int(threadIdx.x) - wid * Gpu::Device::warp_size;
 
        bool sorted = true;
        for (Long i = lane + 1; i < len; i += Gpu::Device::warp_size) {
            sorted = sorted && (pcol[b+i-1] <= pcol[b+i]);
        }
#if defined(AMREX_USE_CUDA)
        if (__all_sync(0xffffffff, sorted)) { return; }
#else
        if (__all(sorted)) { return; }
#endif
 
        constexpr int ITEMS_PER_THREAD = AMREX_HIP_OR_CUDA(2,4);
        constexpr int ITEMS_PER_WARP = Gpu::Device::warp_size * ITEMS_PER_THREAD;
 
        if (len <= ITEMS_PER_WARP)
        {
#if defined(AMREX_USE_CUDA)
            using WarpSort = cub::WarpMergeSort<I, ITEMS_PER_THREAD, Gpu::Device::warp_size, T>;
            __shared__ typename WarpSort::TempStorage temp_storage[nwarps_per_block];
#elif defined(AMREX_USE_HIP)
            using WarpSort = rocprim::warp_sort<I, Gpu::Device::warp_size, T>;
            __shared__ typename WarpSort::storage_type temp_storage[nwarps_per_block];
#endif
 
            I keys[ITEMS_PER_THREAD];
            T values[ITEMS_PER_THREAD];
 
            #pragma unroll
            for (int i = 0; i < ITEMS_PER_THREAD; ++i) {
                int idx = lane * ITEMS_PER_THREAD + i;
                if (idx < len) {
                    keys[i] = pcol[b + idx];
                    values[i] = pmat[b + idx];
                } else {
                    keys[i] = std::numeric_limits<I>::max();
                    values[i] = T{};
                }
            }
 
            AMREX_HIP_OR_CUDA(
                WarpSort{}.sort(keys, values, temp_storage[wid]),
                WarpSort(temp_storage[wid]).Sort(
                    keys, values, [](I x, I y) {return x < y;}));
 
            #pragma unroll
            for (int i = 0; i < ITEMS_PER_THREAD; ++i) {
                int idx = lane * ITEMS_PER_THREAD + i;
                if (idx < len) {
                    pcol[b + idx] = keys[i];
                    pmat[b + idx] = values[i];
                }
            }
        } else {
            if (lane == 0) {
                Gpu::Atomic::AddNoRet(d_needs_fallback, 1);
            }
        }
    });
 
    auto* h_needs_fallback = needs_fallback.copyToHost();
 
    if (*h_needs_fallback)
    {
        // The item count of the segmented sort below is 32-bit.
        AMREX_ALWAYS_ASSERT(nnz < Long(std::numeric_limits<int>::max()));
 
        V<I> col_index_out(col_index.size());
        V<T> mat_out(mat.size());
        auto* d_col_out = col_index_out.data();
        auto* d_val_out = mat_out.data();
 
        std::size_t temp_bytes = 0;
 
        AMREX_GPU_SAFE_CALL(AMREX_HIP_OR_CUDA(
                                rocprim::segmented_radix_sort_pairs,
                                cub::DeviceSegmentedRadixSort::SortPairs)
                            (nullptr, temp_bytes, pcol, d_col_out, pmat, d_val_out,
                             nnz, nr, prow, prow+1, 0, int(sizeof(I)*CHAR_BIT),
                             stream));
 
        auto* d_temp = (void*) The_Arena()->alloc(temp_bytes);
 
        AMREX_GPU_SAFE_CALL(AMREX_HIP_OR_CUDA(
                                rocprim::segmented_radix_sort_pairs,
                                cub::DeviceSegmentedRadixSort::SortPairs)
                            (d_temp, temp_bytes, pcol, d_col_out, pmat, d_val_out,
                             nnz, nr, prow, prow+1, 0, int(sizeof(I)*CHAR_BIT),
                             stream));
 
        std::swap(col_index, col_index_out);
        std::swap(mat, mat_out);
 
        Gpu::streamSynchronize();
        The_Arena()->free(d_temp);
    }
 
    // let's test both by print matrix out to see if it's sorted.
 
    AMREX_GPU_ERROR_CHECK();
 
#elif defined(AMREX_USE_SYCL)
 
    // xxxxx TODO SYCL: Let's not worry about performance for now.
    CSR<T,Gpu::PinnedVector,I> h_csr;
    duplicateCSR(Gpu::deviceToHost, h_csr, *this);
    Gpu::streamSynchronize();
    h_csr.sort_on_host();
    duplicateCSR(Gpu::hostToDevice, *this, h_csr);
    Gpu::streamSynchronize();
 
#endif
 
#else
 
    sort_on_host();
 
#endif
}
 
template <typename T, template <typename> class V, typename I>
void CSR<T,V,I>::sort_on_host ()
{
    if (nnz <= 0) { return; }
 
    constexpr int SMALL = 128;
 
    Long nr = nrows();
 
#ifdef AMREX_USE_OMP
#pragma omp parallel
#endif
    {
        V<I> lcols;
        V<T> lvals;
        V<int> perm;
 
        I scols[SMALL];
        T svals[SMALL];
 
#ifdef AMREX_USE_OMP
#pragma omp for
#endif
        for (Long r = 0; r < nr; ++r) {
            Long const b = row_offset[r  ];
            Long const e = row_offset[r+1];
            auto const len = int(e - b);
 
            if (len <= 1) { continue; }
 
            bool sorted = true;
            for (int i = 1; i < len; ++i) {
                if (col_index[b+i-1] > col_index[b+i]) {
                    sorted = false;
                    break;
                }
            }
            if (sorted) { continue; }
 
            if (len <= SMALL) {
                // Insertion sort using arrays on stack
                for (int i = 0; i < len; ++i) {
                    scols[i] = col_index[b+i];
                    svals[i] = mat      [b+i];
                }
                for (int i = 1; i < len; ++i) {
                    auto c = scols[i];
                    auto v = svals[i];
                    auto j = i;
                    while (j > 0 && scols[j-1] > c) {
                        scols[j] = scols[j-1];
                        svals[j] = svals[j-1];
                        --j;
                    }
                    scols[j] = c;
                    svals[j] = v;
                }
                for (int i = 0; i < len; ++i) {
                    col_index[b+i] = scols[i];
                    mat      [b+i] = svals[i];
                }
            } else {
                lcols.resize(len);
                lvals.resize(len);
                perm.resize(len);
 
                for (int i = 0; i < len; ++i) {
                    lcols[i] = col_index[b+i];
                    lvals[i] = mat      [b+i];
                    perm [i] = i;
                }
 
                std::ranges::sort(perm,
                                  [&] (int i0, int i1) {
                                      return lcols[i0] < lcols[i1];
                                  });
 
                for (int out = 0; out < len; ++out) {
                    auto const in = perm[out];
                    col_index[b+out] = lcols[in];
                    mat      [b+out] = lvals[in];
                }
            }
        }
    }
}
 
// CsrIndex is a helper class for cuSPARSE, because many cuSparse functions
// such as cusparseSpGEMM do not support 64-bit integers.
template <typename I, template <typename> class V>
struct CsrIndex
{
    V<I> col_index;
    V<I> row_offset;
 
    template <typename T, typename J>
    void copyFrom (CsrView<T,J> const& csr)
    {
        AMREX_ALWAYS_ASSERT(csr.nnz < Long(std::numeric_limits<I>::max()) &&
                            csr.nrows < Long(std::numeric_limits<I>::max()));
 
        col_index.resize(csr.nnz);
        row_offset.resize(csr.nrows+1);
 
        auto N = static_cast<I>(std::max(col_index.size(), row_offset.size()));
        auto lnnz = static_cast<I>(col_index.size());
        auto lnro = static_cast<I>(row_offset.size());
        auto* col_dst = col_index.data();
        auto* row_dst = row_offset.data();
        auto const* col_src = csr.col_index;
        auto const* row_src = csr.row_offset;
        ParallelFor(N, [=] AMREX_GPU_DEVICE (I i)
        {
            if (i < lnnz) { col_dst[i] = static_cast<I>(col_src[i]); }
            if (i < lnro) { row_dst[i] = static_cast<I>(row_src[i]); }
        });
    }
 
    template <typename T, typename J>
    void copyTo (CsrView<T,J> const& csr) const
    {
        static_assert(!std::is_const_v<T>);
        auto N = static_cast<I>(std::max(col_index.size(), row_offset.size()));
        auto lnnz = static_cast<I>(col_index.size());
        auto lnro = static_cast<I>(row_offset.size());
        auto const* col_src = col_index.data();
        auto const* row_src = row_offset.data();
        auto* col_dst = csr.col_index;
        auto* row_dst = csr.row_offset;
        ParallelFor(N, [=] AMREX_GPU_DEVICE (I i)
        {
            if (i < lnnz) { col_dst[i] = static_cast<J>(col_src[i]); }
            if (i < lnro) { row_dst[i] = static_cast<J>(row_src[i]); }
        });
    }
 
    template <typename T, template <typename> class U, typename J>
    void copyTo (CSR<T,U,J>& csr) const
    {
        copyTo(csr.view());
    }
};
 
}
 
#endif