matrix.hpp

Go to the documentation of this file.

#ifndef MATRIX_HPP
#define MATRIX_HPP

#include "space_settings.hpp"
#include "my_mirror_view.hpp"

/* The Matrix class stores a sparse matrix. It is initialized in (?) format but can be converted to CSR (Compressed Sparse Row) with the convert_org_to_csr
 * method.
 * */
template<class Device>
class Matrix{

    using exspace = typename Device::execution_space; 
    public: // temporarily public, for testing
    enum SetValueOpt{
        Replace=0,
        Add=1
    };

    enum Type{
        Identity=0
    };

    int m; 
    int n; 
    int width; //< max of non-zero element of each row
    View<double**,CLayout, Device> value; 
    View<int**,CLayout, Device> eindex; 
    View<int*,CLayout, Device> nelement; 

    // use compressed sparse row matrix data structure
    bool is_csr; 
    int nnz; 
    View<int*,CLayout, Device> csr_ridx; 
    View<double*,CLayout, Device> csr_v; 
    View<int*,CLayout, Device> csr_cidx; 

    public:

    Matrix(){}

    Matrix(int m_in,int n_in,int w_in);

    Matrix(Type matrix_type, int m_in);

    Matrix(int m_in, int n_in, int w_in, int nnz_in, bool is_csr_in, int* csr_ridx_or_eindex, int* csr_cidx_or_nelement, double* values);

    // Create a mirror with a different device type
    template<class Device2>
    inline Matrix<Device2> mirror() const{
        Matrix<Device2> mtx;

        mtx.m = m;
        mtx.n = n;
        mtx.width = width;
        mtx.value = my_mirror_view(value, Device2());
        mirror_copy(mtx.value, value);
        mtx.eindex = my_mirror_view(eindex, Device2());
        mirror_copy(mtx.eindex, eindex);
        mtx.nelement = my_mirror_view(nelement, Device2());
        mirror_copy(mtx.nelement, nelement);

        mtx.is_csr = is_csr;
        mtx.nnz = nnz;
        mtx.csr_ridx = my_mirror_view(csr_ridx, Device2());
        mirror_copy(mtx.csr_ridx, csr_ridx);
        mtx.csr_v = my_mirror_view(csr_v, Device2());
        mirror_copy(mtx.csr_v, csr_v);
        mtx.csr_cidx = my_mirror_view(csr_cidx, Device2());
        mirror_copy(mtx.csr_cidx, csr_cidx);
        return mtx;
    }

    void convert_org_to_csr();

    // i row index
    // j column index
    // value_in is value
    // flag is whether to replace or add
    KOKKOS_INLINE_FUNCTION void set_value(int i, int j, double value_in, SetValueOpt flag) const{
        if(is_csr){
            DEVICE_PRINTF("\n Error: set_value is only allowed when matrix is not in csr format");
            return;
        }

        // Check matrix indices are valid
        if(i>=m || j>=n || i<0 || j<0){
            DEVICE_PRINTF("\n Error: Out of bounds access in set_value");
            return;
        }

        // Loop through indices that have already been assigned
        for (int l=0; l<nelement(i); l++){
           if(j==eindex(i,l)-1){ // If this index has already been assigned, replace or add the value (eindex is 1-indexed)
              if(flag==Replace){
                 value(i,l) = value_in;
              } else {
                 value(i,l) += value_in;
              }
              return; //DONE already -- exit the function
           }
        }

        // If the index has not already been assigned, then we need to add the index
        // First, check that we have room in our mapping to add another index
        if(nelement(i) < width){
            int l=nelement(i);
            eindex(i,l)=j+1; // 1-indexed
            value(i,l)=value_in;
            nelement(i) += 1;
        }else{
            DEVICE_PRINTF("\nError in set_value: Not enough memory space for matrix.");
            DEVICE_PRINTF("\nIncrease width. width=%d but nelement(%d)=%d",width,i,nelement(i));
        }
    }

//    private: // Can't be private because there is this cuda error:
//                The enclosing parent function for an extended __host__ __device__ lambda cannot have private or protected access within its class
//                There should be some workarounds available if this is a problem.

    void mult_org(const View<double*,CLayout,Device>& x, const View<double*,Kokkos::LayoutRight,Device>& y) const;

    void mult_csr(const View<double*,CLayout,Device>& x, const View<double*,Kokkos::LayoutRight,Device>& y) const;

    void mult_tensor_org(const View<double**,CLayout,Device>& x, const View<double**,Kokkos::LayoutRight,Device>& y) const;

    void mult_tensor_csr(const View<double**,CLayout,Device>& x, const View<double**,Kokkos::LayoutRight,Device>& y) const;

    void transpose_mult_org(const View<double*,CLayout,Device>& x, const View<double*,Kokkos::LayoutRight,Device>& y) const;

    void transpose_mult_csr(const View<double*,CLayout,Device>& x, const View<double*,Kokkos::LayoutRight,Device>& y) const;

    public:

    void mult(const View<double*,CLayout,Device>& x,const View<double*,Kokkos::LayoutRight,Device>& y) const;

    void mult_tensor(const View<double**,CLayout,Device>& x, const View<double**,Kokkos::LayoutRight,Device>& y) const;

    void transpose_mult(const View<double*,CLayout,Device>& x, const View<double*,Kokkos::LayoutRight,Device>& y) const;

    void write(const std::string stream_name, const std::string filename) const;
};

#endif