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));
        }
    }
 
    KOKKOS_INLINE_FUNCTION void normalize_row(int irow, double norm) const{
        double sum = 0.0;
        for(int i=0; i<nelement(irow); i++) sum += value(irow,i);
        double fac = norm/sum;
        for(int i=0; i<nelement(irow); i++) value(irow,i) *= fac;
    }
 
 
    KOKKOS_INLINE_FUNCTION void sort_row(int i) const{
        for(int j=0; j<nelement(i); j++){
            // Find smallest index
            int ein_min = eindex(i,j);
            double val_min = value(i,j);
            int k=j;
            for(int k2=j;k2<nelement(i);k2++){
                if(eindex(i,k2)<eindex(i,k)){
                    // Save smallest index and value
                    ein_min = eindex(i,k2);
                    val_min = value(i,k2);
                    k = k2;
                }
            }
 
            // Swap values
            eindex(i,k) = eindex(i,j);
            value(i,k) = value(i,j);
            eindex(i,j) = ein_min;
            value(i,j) = val_min;
        }
    }
 
//    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