Line data Source code
1 : !!****m* ABINIT/m_spmat_csr
2 : !! NAME
3 : !! m_spmat_csr
4 : !!
5 : !! FUNCTION
6 : !! This module contains the a CSR (compressed row) format of sparse matrix.
7 : !! Efficient for mat vec multiplication.
8 : !! MPI matvec is also implemented.
9 : !! Datatypes:
10 : !! CSR_mat_t: CSR matrix
11 : !!
12 : !! Subroutines:
13 : !! TODO: add this when F2003 doc style is determined.
14 : !!
15 : !!
16 : !! COPYRIGHT
17 : !! Copyright (C) 2001-2026 ABINIT group (hexu)
18 : !! This file is distributed under the terms of the
19 : !! GNU General Public License, see ~abinit/COPYING
20 : !! or http://www.gnu.org/copyleft/gpl.txt .
21 : !! For the initials of contributors, see ~abinit/doc/developers/contributors.txt .
22 : !!
23 : !! SOURCE
24 :
25 : #if defined HAVE_CONFIG_H
26 : #include "config.h"
27 : #endif
28 : #include "abi_common.h"
29 :
30 : module m_spmat_csr
31 : use defs_basis
32 : use m_abicore
33 : use m_xmpi
34 : use m_errors
35 : use m_spmat_base, only: base_mat2d_t
36 : use m_mpi_scheduler
37 : implicit none
38 : !!***
39 : private
40 : !!----------- CSR ------------------------
41 : ! CSR sparse matrix
42 : ! nnz: number of non-zeros.
43 : ! icol : column index of entries .size:nnz
44 : ! row_shift: row_shift(irow) to row_shift(irow+1)-1 are the index
45 : ! of column and values for entries. size : nrow+1
46 : ! val: values of non-zero entries. size(nnz)
47 : !type, public, extends(base_mat_t) :: CSR_mat_t
48 : type, public, extends(base_mat2d_t) :: CSR_mat_t
49 : integer :: nnz
50 : integer, allocatable :: icol(:), row_shift(:)
51 : real(dp), allocatable:: val(:)
52 : type(mpi_scheduler_t) :: mps
53 : contains
54 : procedure :: initialize
55 : procedure :: set
56 : procedure :: finalize=> csr_mat_t_finalize
57 : procedure :: mv => csr_mat_t_mv ! mat vec multiplication serial version
58 : procedure :: sync ! sync data to all mpi ranks
59 : procedure :: mv_mpi => csr_mat_t_mv_mpi ! mpi version of mv
60 : procedure :: mv_select_row =>csr_mat_t_mv_select_row ! mv of selected rows
61 : procedure :: mv_one_row =>csr_mat_t_mv_one_row ! mv of one rows
62 : procedure :: get_block => csr_mat_t_get_block ! get a block into a dense matrix
63 : end type CSR_mat_t
64 :
65 : contains
66 :
67 : !-----------------------------------------------------------------------
68 : !> @brief initialization
69 : !> @param [in] mshape: shape of matrix, should be size 2.
70 : !-----------------------------------------------------------------------
71 3 : subroutine initialize(self, mshape)
72 : class(csr_mat_t), intent(inout) :: self
73 : integer, intent(in) :: mshape(:)
74 3 : if (size(mshape)/=2) stop "mshape should be size 2"
75 3 : call self%base_mat_t%initialize(mshape)
76 3 : self%nrow=mshape(1)
77 3 : self%ncol=mshape(2)
78 3 : end subroutine initialize
79 :
80 : !-----------------------------------------------------------------------
81 : !> @brief set the full csr matrix
82 : !> @param [in] nnz: number of entries
83 : !> @param [in] icol:
84 : !> @param [in] row_shift:
85 : !> @param [in] val: values:
86 : !-----------------------------------------------------------------------
87 3 : subroutine set(self, nnz, icol, row_shift, val)
88 : class(CSR_mat_t), intent(inout) :: self
89 : integer, intent(in) :: nnz
90 : ! i: col number of each entry
91 : ! j: first 0, n1, n1+n2, ...
92 : ! val(irow, irow+1) are the values of entries in row irow.
93 : integer , intent(in), optional :: icol(:), row_shift(:)
94 : real(dp), intent(in), optional :: val(:)
95 : integer :: iproc
96 :
97 3 : iproc=xmpi_comm_rank(xmpi_world)
98 3 : if (iproc/=0) then
99 0 : ABI_ERROR("This function (CSR_MAT%set) should be only used on root node")
100 : end if
101 :
102 3 : self%nnz=nnz
103 3 : if(.not. allocated(self%icol)) then
104 9 : ABI_MALLOC(self%icol, (self%nnz))
105 : endif
106 :
107 3 : if(.not. allocated(self%row_shift)) then
108 9 : ABI_MALLOC(self%row_shift, (self%nrow+1))
109 : endif
110 :
111 3 : if(.not. allocated(self%val)) then
112 9 : ABI_MALLOC(self%val, (self%nnz))
113 : endif
114 :
115 3 : if (present(icol)) then
116 0 : self%icol(:)=icol(:)
117 : end if
118 3 : if (present(row_shift)) then
119 0 : self%row_shift(:)=row_shift(:)
120 : end if
121 3 : if (present(icol)) then
122 0 : self%val(:)=val(:)
123 : end if
124 :
125 3 : end subroutine set
126 :
127 : !-----------------------------------------------------------------------
128 : !> @brief sync the matrix from master node to all nodes
129 : !> @param [in] master: the id of master node
130 : !> @param [in] comm: the communicator
131 : !> @param [in] nblock: the minimal block for assigning the tasks to nodes.
132 : !-----------------------------------------------------------------------
133 18 : subroutine sync(self, master, comm, nblock)
134 : class (csr_mat_t), intent(inout) :: self
135 : integer , intent(in) :: master, comm, nblock
136 : integer :: ierr, iproc
137 2 : iproc=xmpi_comm_rank(xmpi_world)
138 :
139 2 : call xmpi_barrier(xmpi_world)
140 2 : call xmpi_bcast(self%ndim, master, comm, ierr)
141 2 : call xmpi_bcast(self%ncol, master, comm, ierr)
142 2 : call xmpi_bcast(self%nrow, master, comm, ierr)
143 2 : call xmpi_bcast(self%nnz, master, comm, ierr)
144 :
145 2 : call self%mps%initialize(self%nrow/nblock, master, comm, nblock)
146 2 : if (.not. self%mps%irank==master) then
147 0 : if(.not. allocated(self%mshape)) then
148 0 : ABI_MALLOC(self%mshape, (self%ndim))
149 : endif
150 :
151 0 : if(.not. allocated(self%icol)) then
152 0 : ABI_MALLOC(self%icol, (self%nnz))
153 : endif
154 :
155 0 : if(.not. allocated(self%row_shift)) then
156 0 : ABI_MALLOC(self%row_shift, (self%nrow+1))
157 : endif
158 0 : if(.not. allocated(self%val)) then
159 0 : ABI_MALLOC(self%val, (self%nnz))
160 : endif
161 : end if
162 :
163 : ! TODO: no need to send all the data to all the
164 : ! only the corresponding row data.
165 2 : call xmpi_bcast(self%mshape, master, comm, ierr)
166 2 : call xmpi_bcast(self%icol, master, comm, ierr)
167 2 : call xmpi_bcast(self%row_shift, master, comm, ierr)
168 2 : call xmpi_bcast(self%val, master, comm, ierr)
169 2 : end subroutine sync
170 :
171 :
172 : !-----------------------------------------------------------------------
173 : !> @brief Finalize
174 : !-----------------------------------------------------------------------
175 3 : subroutine CSR_mat_t_finalize(self)
176 : class(CSR_mat_t), intent(inout) :: self
177 3 : self%ncol=0
178 3 : self%nrow=0
179 3 : self%nnz=0
180 3 : self%ndim=0
181 3 : call self%mps%finalize()
182 3 : if(allocated(self%icol)) then
183 3 : ABI_FREE(self%icol)
184 : endif
185 3 : if(allocated(self%row_shift)) then
186 3 : ABI_FREE(self%row_shift)
187 : endif
188 3 : if(allocated(self%val)) then
189 3 : ABI_FREE(self%val)
190 : endif
191 3 : if(allocated(self%mshape)) then
192 3 : ABI_FREE(self%mshape)
193 : endif
194 3 : end subroutine CSR_mat_t_finalize
195 :
196 :
197 : !-----------------------------------------------------------------------
198 : !> @brief Matrix vector multiplication
199 : !> @param [in] x : M x = b
200 : !> @param [out] b: M x = b
201 : !-----------------------------------------------------------------------
202 42000 : subroutine CSR_mat_t_mv(self, x, b)
203 : class(CSR_mat_t), intent(in):: self
204 : real(dp), intent(in) :: x(self%ncol)
205 : real(dp), intent(out) :: b(self%nrow)
206 : integer::irow, i1, i2, i
207 43050000 : b(:)=0.0d0
208 : !$OMP PARALLEL DO private(i, i1, i2)
209 43050000 : do irow=1, self%nrow
210 43008000 : i1=self%row_shift(irow)
211 43008000 : i2=self%row_shift(irow+1)-1
212 2881578000 : do i=i1, i2
213 2881536000 : b(irow)=b(irow)+ self%val(i)*x(self%icol(i))
214 : end do
215 : enddo
216 : !$OMP END PARALLEL DO
217 42000 : end subroutine CSR_mat_t_mv
218 :
219 : !-----------------------------------------------------------------------
220 : !> @brief Matrix vector multiplication (mpi version)
221 : !> @param [in] x : M x = b
222 : !> @param [out] b: M x = b
223 : !-----------------------------------------------------------------------
224 4006 : subroutine CSR_mat_t_mv_mpi(self, x, b, bcastx, syncb)
225 : class(CSR_mat_t), intent(in) :: self
226 : real(dp), intent(inout) :: x(self%ncol)
227 : logical, intent(in) :: bcastx, syncb
228 : real(dp), intent(out) :: b(self%nrow)
229 : !real(dp):: my_b(self%nrow)
230 : integer :: ierr, irow, i1, i2, i
231 4006 : if (bcastx) then
232 0 : call xmpi_bcast(x, 0, xmpi_world, ierr)
233 : end if
234 : !if (.not. iam_master) b(:)=0.0_dp
235 : !my_b(:)=0.0_dp
236 2599894 : b(:)=0.0_dp
237 2599894 : do irow= self%mps%istart, self%mps%iend
238 2595888 : i1=self%row_shift(irow)
239 2595888 : i2=self%row_shift(irow+1)-1
240 70092982 : do i=i1, i2
241 70088976 : b(irow)=b(irow)+ self%val(i)*x(self%icol(i))
242 : end do
243 : enddo
244 : ! TODO : use gather instead of reduce?
245 : !call mpi_reduce(my_b, b, self%nrow, MPI_DOUBLE_PRECISION, MPI_SUM, 0, MPI_COMM_WORLD, ierr)
246 4006 : if (syncb) then
247 0 : call xmpi_sum_master(b, 0, xmpi_world, ierr )
248 : endif
249 4006 : end subroutine CSR_mat_t_mv_mpi
250 :
251 : !-----------------------------------------------------------------------
252 : !> @brief multiple a row, indexed by of M by x: y_i= Mij x_j
253 : !> @param [in] nrow: number of rows.
254 : !> @param [in] ind_row: indices i
255 : !> @param [in] x: x
256 : !> @param [out] y: y
257 : !-----------------------------------------------------------------------
258 0 : subroutine CSR_mat_t_mv_one_row(self, j, x, y)
259 : class(CSR_mat_t), intent(in)::self
260 : integer, intent(in) ::j
261 : real(dp), intent(in) :: x(self%ncol)
262 : real(dp), intent(out) :: y
263 : integer :: i, i1, i2
264 0 : y=0.0_dp
265 0 : i1=self%row_shift(j)
266 0 : i2=self%row_shift(j+1)-1
267 0 : do i=i1, i2
268 0 : y= y+self%val(i)*x(self%icol(i))
269 : end do
270 0 : end subroutine CSR_mat_t_mv_one_row
271 :
272 :
273 : !-----------------------------------------------------------------------
274 : !> @brief multiple a submatrix (rows indexed by i) of M by x: y_i=\sum M_ij x_j
275 : !> @param [in] nrow: number of rows.
276 : !> @param [in] ind_row: indices i
277 : !> @param [in] x: x
278 : !> @param [out] y: y
279 : !-----------------------------------------------------------------------
280 0 : subroutine CSR_mat_t_mv_select_row(self, nrow, id_row, x, y)
281 : class(CSR_mat_t), intent(in)::self
282 : integer, intent(in) :: nrow, id_row(nrow)
283 : real(dp), intent(in) :: x(self%ncol)
284 : real(dp), intent(out) :: y(nrow)
285 : integer :: i, irow, i1, i2, j
286 0 : y(:)=0.0_dp
287 0 : do j=1, nrow
288 0 : irow=id_row(j)
289 0 : i1=self%row_shift(irow)
290 0 : i2=self%row_shift(irow+1)-1
291 0 : do i=i1, i2
292 0 : y(j)= y(j)+self%val(i)*x(self%icol(i))
293 : end do
294 : end do
295 0 : end subroutine CSR_mat_t_mv_select_row
296 :
297 432 : subroutine CSR_mat_t_get_block(self,irow_start, icol_start, nrow, ncol, blk)
298 : class(CSR_mat_t), intent(in)::self
299 : integer, intent(in) :: nrow, irow_start, ncol, icol_start
300 : real(dp), intent(out) :: blk(nrow, ncol)
301 : integer :: i, irow, i1, i2, j, icol
302 5616 : blk(:, :)=0.0_dp
303 1728 : do j=1, nrow
304 1296 : irow = irow_start+j-1
305 1296 : i1=self%row_shift(irow)
306 1296 : i2=self%row_shift(irow+1)-1
307 35424 : do i=i1, i2
308 33696 : icol = self%icol(i)
309 34992 : if(icol>=icol_start .and. icol<icol_start+ncol) then
310 0 : blk(j, icol-icol_start+1) = self%val(i)
311 : endif
312 : end do
313 : end do
314 432 : end subroutine CSR_mat_t_get_block
315 :
316 : ! subroutine print(self)
317 : ! class(CSR_mat_t), intent(in) :: self
318 : ! print *, "icol:", self%icol
319 : ! print *, "row_shift:", self%row_shift
320 : ! print *, "val:", self%val
321 : ! end subroutine print
322 :
323 0 : end module m_spmat_csr
|
