Skip to content

Commit

Permalink
SpMatrix: Add constructor for CSR format
Browse files Browse the repository at this point in the history
  • Loading branch information
WeiqunZhang committed Jan 28, 2025
1 parent 2f3a8c7 commit ca18578
Showing 1 changed file with 30 additions and 0 deletions.
30 changes: 30 additions & 0 deletions Src/LinearSolvers/AMReX_SpMatrix.H
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,14 @@ public:

void define (AlgPartition partition, int nnz);

//! Define a matrix with CSR format data. Note that mat and col_index
//! should contains nnz elements. The number of elements in row_index
//! should the numbrer of local rows plus 1. The data can be freed after
//! this function call. For GPU builds, the data are expected to be in
//! GPU memory.
void define (AlgPartition partition, T const* mat, Long const* col_index,
Long nnz, Long const* row_index);

[[nodiscard]] AlgPartition const& partition () const { return m_partition; }

[[nodiscard]] Long numLocalRows () const { return m_row_end - m_row_begin; }
Expand Down Expand Up @@ -160,6 +168,28 @@ SpMatrix<T,Allocator>::define_doit (int nnz)
});
}

template <typename T, template<typename> class Allocator>
void
SpMatrix<T,Allocator>::define (AlgPartition partition, T const* mat,
Long const* col_index, Long nnz,
Long const* row_index)
{
m_partition = std::move(partition);
m_row_begin = m_partition[ParallelDescriptor::MyProc()];
m_row_end = m_partition[ParallelDescriptor::MyProc()+1];
Long nlocalrows = this->numLocalRows();
m_data.mat.resize(nnz);
m_data.col_index.resize(nnz);
m_data.row_offset.resize(nlocalrows+1);
m_data.nnz = nnz;
Gpu::copyAsync(Gpu::deviceToDevice, mat, mat+nnz, m_data.mat.begin());
Gpu::copyAsync(Gpu::deviceToDevice, col_index, col_index+nnz,
m_data.col_index.begin());
Gpu::copyAsync(Gpu::deviceToDevice, row_index, row_index+nlocalrows+1,
m_data.row_index.begin());
Gpu::streamSynchronize();
}

template <typename T, template<typename> class Allocator>
void
SpMatrix<T,Allocator>::printToFile (std::string const& file) const
Expand Down

0 comments on commit ca18578

Please sign in to comment.