gradient_matrices.hpp

Go to the documentation of this file.

#ifndef GRADIENT_MATRICES_HPP
#define GRADIENT_MATRICES_HPP

#include "matrix.hpp"

extern "C" void get_grad_mat_data(int ind, int* m, int* n, int* w, int* nnz, int* is_csr_int, int** csr_ridx_or_eindex, int** csr_cidx_or_nelement, double** values);

template<class Device>
struct GradientMatrices{
    Matrix<Device> gradx;
    Matrix<Device> grady;

    GradientMatrices(){}

    /* Set up gradx and grady on GPU by (for now) copying from the fortran matrix */
    GradientMatrices(bool use){
        int m, n, w, nnz, is_csr_int;
        int* csr_ridx_or_eindex;
        int* csr_cidx_or_nelement;
        double* values;
        const int X=0;
        const int Y=1;
        // gradx
        get_grad_mat_data(X, &m, &n, &w, &nnz, &is_csr_int, &csr_ridx_or_eindex, &csr_cidx_or_nelement, &values);
        gradx = Matrix<Device>(m, n, w, nnz, (is_csr_int==1), csr_ridx_or_eindex, csr_cidx_or_nelement, values);

        // grady
        get_grad_mat_data(Y, &m, &n, &w, &nnz, &is_csr_int, &csr_ridx_or_eindex, &csr_cidx_or_nelement, &values);
        grady = Matrix<Device>(m, n, w, nnz, (is_csr_int==1), csr_ridx_or_eindex, csr_cidx_or_nelement, values);
    }

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

        m.gradx = gradx.template mirror<Device2>();
        m.grady = grady.template mirror<Device2>();

        return m;
    }

    GradientMatrices(int nnode, int matrix_width)
      : gradx(nnode, nnode, matrix_width),
        grady(nnode, nnode, matrix_width){}

};
#endif