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_h_view_t = Kokkos::View<value_type****, Kokkos::LayoutRight, host_type>;
    using const_vector_h_view_t = Kokkos::View<const value_type****, Kokkos::LayoutRight, host_type>;
    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_h_view_t b_h, vector_h_view_t x_h, const_vector_view_t b, vector_view_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_h_view_t = typename GridMatrix<Device>::vector_h_view_t;
    using const_vector_h_view_t = typename GridMatrix<Device>::const_vector_h_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_h_view_t b_h, vector_h_view_t x_h, const_vector_view_t b, vector_view_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_h_view_t = typename GridMatrix<Device>::vector_h_view_t;
    using const_vector_h_view_t = typename GridMatrix<Device>::const_vector_h_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_h_view_t b_h, vector_h_view_t x_h, const_vector_view_t b, vector_view_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
};
 
 
std::unique_ptr<GridMatrix<DeviceType>> create_matrix(int n_matrices,
    int nvr, int nvz, int num_species, LinAlgBackend labackend,
    double ginkgo_residual_reduction, int ginkgo_max_iterations,
    std::string format = "default");
 
 
}
 
#endif