matrix.hpp

Go to the documentation of this file.

#ifndef MATRIX_HPP
#define MATRIX_HPP

#include "space_settings.hpp"
#include "array_deep_copy.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; 

    enum SetValueOpt{
        Replace=0,
        Add=1
    };

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

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

    public:

    Matrix(){}

    Matrix(int m_in,int n_in,int w_in)
      : m(m_in),
        n(n_in),
        width(w_in),
        value("value",m_in,w_in), // init to zero
        eindex(Kokkos::ViewAllocateWithoutInitializing("eindex"), m_in,w_in),
        nelement("nelement",m_in), // init to zero
        is_csr(false) // start from fixed width matrix format
      {}

    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)
      : m(m_in),
        n(n_in),
        width(w_in),
        nnz(nnz_in),
        is_csr(is_csr_in)
    {
        if(is_csr){
            csr_ridx = Kokkos::View<int*,Kokkos::LayoutRight, Device>(Kokkos::ViewAllocateWithoutInitializing("csr_ridx"),m+1);
            csr_v = Kokkos::View<double*,Kokkos::LayoutRight, Device>(Kokkos::ViewAllocateWithoutInitializing("csr_v"),nnz);
            csr_cidx = Kokkos::View<int*,Kokkos::LayoutRight, Device>(Kokkos::ViewAllocateWithoutInitializing("csr_cidx"),nnz);
            array_deep_copy(csr_ridx, csr_ridx_or_eindex);
            array_deep_copy(csr_v, values);
            array_deep_copy(csr_cidx, csr_cidx_or_nelement);
        } else {
            value = Kokkos::View<double**,Kokkos::LayoutRight, Device>(Kokkos::ViewAllocateWithoutInitializing("value"), m,width);
            eindex = Kokkos::View<int**,Kokkos::LayoutRight, Device>(Kokkos::ViewAllocateWithoutInitializing("eindex"), m,width);
            nelement = Kokkos::View<int*,Kokkos::LayoutRight, Device>(Kokkos::ViewAllocateWithoutInitializing("nelement"), m);
            array_deep_copy(eindex, csr_ridx_or_eindex);
            array_deep_copy(value, values);
            array_deep_copy(nelement, csr_cidx_or_nelement);
        }
    }

/* NOT READY YET - NEEDS TO CHECK CPU V GPU
    **
     * Converts matrix from original format to CSR.
     *
     * @return void
     *
    void convert_org_to_csr(){
        if(is_csr){
            printf("\nwarning:matrix is csr already\n");
            return;
        }

        //counting nonzero elements
        is_csr=true;
        
        nnz = 0;
        for (int i=0; i<m; i++){
            nnz+=nelement(i);
        }

        //allocate memory -- 0-based index for C++ compatiblity
        csr_ridx = Kokkos::View<int*,Kokkos::LayoutRight, Device>(Kokkos::ViewAllocateWithoutInitializing("csr_ridx"),m+1);
        csr_v = Kokkos::View<double*,Kokkos::LayoutRight, Device>(Kokkos::ViewAllocateWithoutInitializing("csr_v"),nnz);
        csr_cidx = Kokkos::View<int*,Kokkos::LayoutRight, Device>(Kokkos::ViewAllocateWithoutInitializing("csr_cidx"),nnz);

        //sorting of eindex is required if element searching is required.
        //currently unsorted -- maybe soring can speed up matrix multiplication?

        //assign values
        csr_ridx(0)=0; // start from 0 - 0 based
        for (int i=0; i<m; i++){
          csr_ridx(i+1)=csr_ridx(i)+nelement(i);
          for (int j=0; j<nelement(i); j++){
            int loc=csr_ridx(i)+j;
            csr_cidx(loc)=eindex(i,j);
            csr_v(loc)=value(i,j);
          }
        }

        //deallocate org data structure
        value = Kokkos::View<double**,Kokkos::LayoutRight, Device>;
        eindex = Kokkos::View<int**,Kokkos::LayoutRight, Device>;
        nelement = Kokkos::View<int*,Kokkos::LayoutRight, Device>;
    }

    // i row index
    // j column index
    // value_in is value
    // flag is whether to replace or add
    **
     * Set value of matrix component. Inlined since this will presumably be called a lot in a loop
     *
     * @param[in] i is the row index
     * @param[in] j is the column index
     * @param[in] value_in is the value to set
     * @param[in] flag is whether to add or replace if there is an existing value in the chosen matrix element
     * @return void
     *
    inline void set_value(int i, int j, double value_in, SetValueOpt flag){
        if(is_csr){
          exit_XGC("\n Error: set_value is only allowed when matrix is not in csr format");
        }

        // error check
        if(i>=m || j>=n || i<0 || j<0){
          exit_XGC("\n Error: Out of bounds access in set_value");
        }

        // check same node
        for (int l=0; l<nelement(i); l++){
           // redundant point
           if(j==eindex(i,l)){
              if(flag==Replace){
                 value(i,l) = value_in;
              } else {
                 value(i,l) += value_in;
              }
              return; //DONE already -- exit the function
           }
        }

        // new point
        nelement(i) += 1;
        int l=nelement(i);

        // Check l is within bounds
        if(l < width){
           printf("Error: Not enough memory space for matrix.");
           printf("Increase width. %d  %d  %d",width,nelement(i),i;
           exit_XGC("\n Error in set_value");
        }

        eindex(i,l)=j;
        value(i,l)=value;
    }
*/
//    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 Kokkos::View<double*,Kokkos::LayoutRight,Device>& x, const Kokkos::View<double*,Kokkos::LayoutRight,Device>& y) const{
        Matrix<Device> mat = *this; // Make local shallow copy so class members are captured by lambda
        Kokkos::parallel_for("mult_org", Kokkos::RangePolicy<exspace>(0,m), KOKKOS_LAMBDA( const int i ){
            y(i)=0.0;
            for (int j=0; j<mat.nelement(i); j++){
                y(i) += mat.value(i,j)*x(mat.eindex(i,j)-1);
            }
        });
    }

    void mult_csr(const Kokkos::View<double*,Kokkos::LayoutRight,Device>& x, const Kokkos::View<double*,Kokkos::LayoutRight,Device>& y) const{
        Matrix<Device> mat = *this; // Make local shallow copy so class members are captured by lambda
        Kokkos::parallel_for("mult_csr", Kokkos::RangePolicy<exspace>(0,m), KOKKOS_LAMBDA( const int i ){
            y(i)=0.0;
            for (int j=mat.csr_ridx(i); j<mat.csr_ridx(i+1); j++){
                y(i) += mat.csr_v(j)*x(mat.csr_cidx(j));
            }
        });
    }

/* Not ready - may need atomics
    **
     * Matrix multiplication using the transpose, i.e.: y = A^T x. Matrix is in original data format.
     *
     * @param[in] x is the input vector
     * @param[out] y is the result
     * @return void
     *
    void transpose_mult_org(const Kokkos::View<double*,Kokkos::LayoutRight,Device>& x, Kokkos::View<double*,Kokkos::LayoutRight,Device>& y){
        Kokkos::deep_copy(y,0.0);
        for (i=0; i<m; i++){
            for (j=0; j<nelement(i); j++){
                int k=eindex(i,j)-1;
                y(k) += value(i,j)*x(i);
            }
        });
    }

    **
     * Matrix multiplication using the transpose, i.e.: y = A^T x. Matrix is in CSR format
     *
     * @param[in] x is the input vector 
     * @param[out] y is the result
     * @return void
     *
    void transpose_mult_csr(const Kokkos::View<double*,Kokkos::LayoutRight,Device>& x, Kokkos::View<double*,Kokkos::LayoutRight,Device>& y){
        Kokkos::deep_copy(y,0.0);
        for (i=0; i<m; i++){
            for (j=csr_ridx(i); j<csr_ridx(i+1); j++){
                int k=csr_cidx(j);
                y(k) += csr_v(j)*x(i);
            }
        });
    }
*/

    void mult_tensor_org(const Kokkos::View<double**,Kokkos::LayoutRight,Device>& x, const Kokkos::View<double**,Kokkos::LayoutRight,Device>& y) const{
        int nv = x.extent(1); // width of tensor
        Matrix<Device> mat = *this; // Make local shallow copy so class members are captured by lambda
        Kokkos::parallel_for("mult_tensor_org", Kokkos::RangePolicy<exspace>(0,m*nv), KOKKOS_LAMBDA( const int idx ){
            int k = idx%nv;
            int i = idx/nv;
            y(i,k)=0.0;
            for (int j=0; j<=mat.nelement(i); j++){
                y(i,k) += mat.value(i,j)*x(mat.eindex(i,j)-1,k);
            }
        });
    }

    void mult_tensor_csr(const Kokkos::View<double**,Kokkos::LayoutRight,Device>& x, const Kokkos::View<double**,Kokkos::LayoutRight,Device>& y) const{
        int nv = x.extent(1); // width of tensor
        Matrix<Device> mat = *this; // Make local shallow copy so class members are captured by lambda
        Kokkos::parallel_for("mult_tensor_csr", Kokkos::RangePolicy<exspace>(0,m*nv), KOKKOS_LAMBDA( const int idx ){
            int k = idx%nv;
            int i = idx/nv;
            y(i,k)=0.0;
            for (int j=mat.csr_ridx(i); j<mat.csr_ridx(i+1); j++){
                y(i,k) += mat.csr_v(j)*x(mat.csr_cidx(j),k);
            }
        });
    }

    public:

    void mult(const Kokkos::View<double*,Kokkos::LayoutRight,Device>& x,const Kokkos::View<double*,Kokkos::LayoutRight,Device>& y) const{
        if(is_csr){
            mult_csr(x,y);
        } else {
            mult_org(x,y);
        }
    }

/*    void transpose_mult(const Kokkos::View<double*,Kokkos::LayoutRight,Device>& x, Kokkos::View<double*,Kokkos::LayoutRight,Device>& y){
        if(is_csr){
            transpose_mult_csr(x,y);
        } else {
            transpose_mult_org(x,y);
        }
    }*/

    void mult_tensor(const Kokkos::View<double**,Kokkos::LayoutRight,Device>& x, const Kokkos::View<double**,Kokkos::LayoutRight,Device>& y) const{
        if(is_csr){
            mult_tensor_csr(x,y);
        } else {
            mult_tensor_org(x,y);
        }
    }
};

#endif