col_grid_matrix.hpp

Go to the documentation of this file.

#ifndef XGC_CORE_CPP_COL_GRID_MATRIX_HPP_
#define XGC_CORE_CPP_COL_GRID_MATRIX_HPP_

#include <stdexcept>
#include <memory>

#include <Kokkos_DualView.hpp>

#include "space_settings.hpp"

#ifdef USE_GINKGO
#include <ginkgo/ginkgo.hpp>
#endif


namespace Collisions {


using size_type = int;


template <typename Device>
class GridMatrix
{
public:
    using index_type = int;
    using host_type = HostType; //Kokkos::HostSpace;
    using device = Device;
    using value_type = double;
    using values_array_t = Kokkos::View<value_type***, Kokkos::LayoutRight, Device,
                                            Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
    using const_values_array_t =
        Kokkos::View<const value_type***, Kokkos::LayoutRight, Device,
                     Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
    using vector_view_t = Kokkos::View<value_type****, Kokkos::LayoutRight, Device>;
    using const_vector_view_t = Kokkos::View<const value_type****, Kokkos::LayoutRight, Device>;
    using vector_dualview_t = Kokkos::DualView<value_type****, Kokkos::LayoutRight, Device>;
    using const_vector_dualview_t = Kokkos::DualView<const value_type****, Kokkos::LayoutRight, Device>;

    GridMatrix(size_type n_nodes, int n_species)
        : num_nodes_{n_nodes}, num_species_{n_species}, batch_size_{n_nodes * n_species}
    { }

    virtual ~GridMatrix() { }

    typename Kokkos::View<index_type**, device>::HostMirror host_index_map_LU() {
        return index_map_.h_view;
    }

    Kokkos::View<index_type**, device> device_index_map_LU() const {
        return index_map_.d_view;
    }

    virtual const_values_array_t get_const_values() const = 0;

    virtual values_array_t get_values() = 0;

    virtual void subtract_from_identity() = 0;

    // For picard_step_init
    virtual void add_identity_multiply(const_vector_view_t dist_col,
                                       vector_view_t dist_iter) const = 0;

    virtual void apply_solve(const_vector_dualview_t b, vector_dualview_t x) const = 0;

    index_type get_num_nonzeros() const { return nnz_; }
    index_type get_num_rows() const { return num_rows_; }

protected:
    Kokkos::DualView<index_type**,device> index_map_;
    const index_type num_species_;
    const index_type num_nodes_;
    const size_type batch_size_;
    index_type num_rows_;
    index_type nnz_;
};


template <typename Device>
class CSCMatrix : public GridMatrix<Device>
{
public:
    using index_type = typename GridMatrix<Device>::index_type;
    using value_type = typename GridMatrix<Device>::value_type;
    using host_type = typename GridMatrix<Device>::host_type;
    using device = typename GridMatrix<Device>::device;
    using values_array_t = typename GridMatrix<Device>::values_array_t;
    using const_values_array_t = typename GridMatrix<Device>::const_values_array_t;
    using internal_values_array_t = Kokkos::DualView<value_type***, Kokkos::LayoutRight, Device>;
    using vector_view_t = typename GridMatrix<Device>::vector_view_t;
    using const_vector_view_t = typename GridMatrix<Device>::const_vector_view_t;
    using vector_dualview_t = typename GridMatrix<Device>::vector_dualview_t;
    using const_vector_dualview_t = typename GridMatrix<Device>::const_vector_dualview_t;

    CSCMatrix(size_type n_nodes, index_type num_species, int nvr, int nvz)
        : GridMatrix<Device>(n_nodes, num_species), nvr_{nvr},
        nvz_{nvz}
    {
        setup_sparsity();
    }

    const_values_array_t get_const_values() const override;

    values_array_t get_values() override;

    void subtract_from_identity() override;

    void add_identity_multiply(const_vector_view_t dist_col,
                               vector_view_t dist_iter) const override;

    void apply_solve(const_vector_dualview_t b, vector_dualview_t x) const override;

protected:
    using GridMatrix<Device>::index_map_;
    using GridMatrix<Device>::num_nodes_;
    using GridMatrix<Device>::num_rows_;
    using GridMatrix<Device>::nnz_;
    using GridMatrix<Device>::batch_size_;
    using GridMatrix<Device>::num_species_;

    index_type nvz_;
    index_type nvr_;

    internal_values_array_t values_;
    Kokkos::View<index_type*,device> central_locs_;
    Kokkos::View<int*,host_type> LU_rowindx_;
    Kokkos::View<int*,host_type> LU_colptr_;

    void setup_sparsity();
};

#ifdef USE_GINKGO

template <typename Device>
class ELLMatrix : public GridMatrix<Device>
{
public:
    using index_type = typename GridMatrix<Device>::index_type;
    using host_type = typename GridMatrix<Device>::host_type;
    using device = typename GridMatrix<Device>::device;
    using value_type = typename GridMatrix<Device>::value_type;
    using values_array_t = typename GridMatrix<Device>::values_array_t;
    using const_values_array_t = typename GridMatrix<Device>::const_values_array_t;
    using vector_view_t = typename GridMatrix<Device>::vector_view_t;
    using const_vector_view_t = typename GridMatrix<Device>::const_vector_view_t;
    using vector_dualview_t = typename GridMatrix<Device>::vector_dualview_t;
    using const_vector_dualview_t = typename GridMatrix<Device>::const_vector_dualview_t;
    using mtx_type = gko::batch::matrix::Ell<value_type, index_type>;

    static constexpr int max_nnz_per_row = 9;

    ELLMatrix(std::shared_ptr<const gko::Executor> exec, size_type n_matrices,
              index_type num_species, int nvr, int nvz,
              double ginkgo_residual_reduction, int ginkgo_max_iterations)
        : exec_{exec}, GridMatrix<Device>(n_matrices, num_species), nvr_{nvr},
          nvz_{nvz}, ginkgo_residual_reduction_{ginkgo_residual_reduction},
        ginkgo_max_iterations_{ginkgo_max_iterations}
    {
        setup_sparsity();
    }

    const_values_array_t get_const_values() const override;

    values_array_t get_values() override;

    void subtract_from_identity() override;

    void add_identity_multiply(const_vector_view_t dist_col,
                               vector_view_t dist_iter) const override;

    void apply_solve(const_vector_dualview_t b, vector_dualview_t x) const override;

protected:
    using GridMatrix<Device>::index_map_;
    using GridMatrix<Device>::num_nodes_;
    using GridMatrix<Device>::num_rows_;
    using GridMatrix<Device>::nnz_;
    using GridMatrix<Device>::batch_size_;
    using GridMatrix<Device>::num_species_;
    std::shared_ptr<const gko::Executor> exec_;
    std::shared_ptr<mtx_type> mtx_;
    
    index_type nvz_;
    index_type nvr_;

    double ginkgo_residual_reduction_;
    int ginkgo_max_iterations_;
  
    void setup_sparsity();
};

#endif

template <typename Device>
class InvalidMatrixType : public std::exception
{
public:
    InvalidMatrixType(const GridMatrix<Device>* const mtx)
    {
        std::string mtx_type = "unknown";
        if(dynamic_cast<const CSCMatrix<Device>*>(mtx)) {
            mtx_type = "CSC";
        }
        what_ = mtx_type + " matrix type not supported!";
    }

    virtual const char* what() const noexcept override { return what_.c_str(); }

private:
    std::string what_;
};


enum class LinAlgBackend {
    lapack,
    ginkgo
};


template <typename Device>
std::unique_ptr<GridMatrix<Device>> create_matrix(size_type n_matrices,
    int nvr, int nvz, int num_species, LinAlgBackend labackend,
    double ginkgo_residual_reduction, int ginkgo_max_iterations,
    std::string format = "default");


}


#include "col_grid_matrix.tpp"


#endif