New SPSPMM implementation

This commit is contained in:
Salvatore Filippone
2023-08-29 19:27:29 +02:00
parent d0cacda995
commit 0d8a5d3dc2
15 changed files with 3271 additions and 168 deletions
+262 -138
View File
@@ -3213,10 +3213,10 @@ subroutine psb_d_cp_csr_from_coo(a,b,info)
if (info == psb_success_) call psb_safe_ab_cpy(b%ja,a%ja,info)
if (info == psb_success_) call psb_safe_ab_cpy(b%val,a%val,info)
if (info == psb_success_) call psb_realloc(max(nr+1,nc+1),a%irp,info)
endif
#if defined(OPENMP)
!$OMP PARALLEL default(shared) reduction(max:info)
@@ -3246,7 +3246,7 @@ subroutine psb_d_cp_csr_from_coo(a,b,info)
#endif
call a%set_host()
end subroutine psb_d_cp_csr_from_coo
@@ -3495,7 +3495,7 @@ subroutine psb_d_cp_csr_to_fmt(a,b,info)
if (info == 0) call psb_safe_cpy( a%ja(1:nz), b%ja , info)
if (info == 0) call psb_safe_cpy( a%val(1:nz), b%val , info)
else
! Despite the implementation in safe_cpy, it seems better this way
! Despite the implementation in safe_cpy, it seems better this way
call psb_realloc(nr+1,b%irp,info)
call psb_realloc(nz,b%ja,info)
call psb_realloc(nz,b%val,info)
@@ -3600,7 +3600,7 @@ subroutine psb_d_cp_csr_from_fmt(a,b,info)
if (info == 0) call psb_safe_cpy( b%ja(1:nz) , a%ja , info)
if (info == 0) call psb_safe_cpy( b%val(1:nz) , a%val , info)
else
! Despite the implementation in safe_cpy, it seems better this way
! Despite the implementation in safe_cpy, it seems better this way
call psb_realloc(nr+1,a%irp,info)
call psb_realloc(nz,a%ja,info)
call psb_realloc(nz,a%val,info)
@@ -3654,6 +3654,130 @@ subroutine psb_d_csr_clean_zeros(a, info)
call a%set_host()
end subroutine psb_d_csr_clean_zeros
#if 0
subroutine psb_dcsrspspmm(a,b,c,info)
use psb_d_mat_mod
use psb_serial_mod, psb_protect_name => psb_dcsrspspmm
implicit none
class(psb_d_csr_sparse_mat), intent(in) :: a,b
type(psb_d_csr_sparse_mat), intent(out) :: c
integer(psb_ipk_), intent(out) :: info
integer(psb_ipk_) :: ma,na,mb,nb, nzc, nza, nzb
character(len=20) :: name
integer(psb_ipk_) :: err_act
name='psb_csrspspmm'
call psb_erractionsave(err_act)
info = psb_success_
if (a%is_dev()) call a%sync()
if (b%is_dev()) call b%sync()
ma = a%get_nrows()
na = a%get_ncols()
mb = b%get_nrows()
nb = b%get_ncols()
if ( mb /= na ) then
write(psb_err_unit,*) 'Mismatch in SPSPMM: ',ma,na,mb,nb
info = psb_err_invalid_matrix_sizes_
call psb_errpush(info,name)
goto 9999
endif
! Estimate number of nonzeros on output.
nza = a%get_nzeros()
nzb = b%get_nzeros()
nzc = 2*(nza+nzb)
call c%allocate(ma,nb,nzc)
call csr_spspmm(a,b,c,info)
call c%set_asb()
call c%set_host()
call psb_erractionrestore(err_act)
return
9999 call psb_error_handler(err_act)
return
contains
subroutine csr_spspmm(a,b,c,info)
implicit none
type(psb_d_csr_sparse_mat), intent(in) :: a,b
type(psb_d_csr_sparse_mat), intent(inout) :: c
integer(psb_ipk_), intent(out) :: info
integer(psb_ipk_) :: ma,na,mb,nb
integer(psb_ipk_), allocatable :: irow(:), idxs(:)
real(psb_dpk_), allocatable :: row(:)
integer(psb_ipk_) :: i,j,k,irw,icl,icf, iret, &
& nzc,nnzre, isz, ipb, irwsz, nrc, nze
real(psb_dpk_) :: cfb
info = psb_success_
ma = a%get_nrows()
na = a%get_ncols()
mb = b%get_nrows()
nb = b%get_ncols()
nze = min(size(c%val),size(c%ja))
isz = max(ma,na,mb,nb)
call psb_realloc(isz,row,info)
if (info == 0) call psb_realloc(isz,idxs,info)
if (info == 0) call psb_realloc(isz,irow,info)
if (info /= 0) return
row = dzero
irow = 0
nzc = 1
do j = 1,ma
c%irp(j) = nzc
nrc = 0
do k = a%irp(j), a%irp(j+1)-1
irw = a%ja(k)
cfb = a%val(k)
irwsz = b%irp(irw+1)-b%irp(irw)
do i = b%irp(irw),b%irp(irw+1)-1
icl = b%ja(i)
if (irow(icl)<j) then
nrc = nrc + 1
idxs(nrc) = icl
irow(icl) = j
end if
row(icl) = row(icl) + cfb*b%val(i)
end do
end do
if (nrc > 0 ) then
if ((nzc+nrc)>nze) then
nze = max(ma*((nzc+j-1)/j),nzc+2*nrc)
call psb_realloc(nze,c%val,info)
if (info == 0) call psb_realloc(nze,c%ja,info)
if (info /= 0) return
end if
call psb_qsort(idxs(1:nrc))
do i=1, nrc
irw = idxs(i)
c%ja(nzc) = irw
c%val(nzc) = row(irw)
row(irw) = dzero
nzc = nzc + 1
end do
end if
end do
c%irp(ma+1) = nzc
end subroutine csr_spspmm
end subroutine psb_dcsrspspmm
#else
subroutine psb_dcsrspspmm(a,b,c,info)
use psb_d_mat_mod
use psb_serial_mod, psb_protect_name => psb_dcsrspspmm
@@ -3791,7 +3915,7 @@ contains
end subroutine csr_spspmm
! gustavson's algorithm using perfect hashing
! and OpenMP parallelisation
! and OpenMP parallelisation
subroutine spmm_omp_gustavson(a,b,c,info)
use omp_lib
@@ -3799,12 +3923,12 @@ contains
type(psb_d_csr_sparse_mat), intent(in) :: a,b
type(psb_d_csr_sparse_mat), intent(out):: c
integer(psb_ipk_), intent(out) :: info
real(psb_dpk_), allocatable :: vals(:), acc(:)
integer(psb_ipk_) :: ma, nb
integer(psb_ipk_), allocatable :: col_inds(:), offsets(:)
integer(psb_ipk_) :: irw, jj, j, k, nnz, rwnz, thread_upperbound, start_idx, end_idx
ma = a%get_nrows()
nb = b%get_ncols()
@@ -3814,7 +3938,7 @@ contains
! dense accumulator
! https://sc18.supercomputing.org/proceedings/workshops/workshop_files/ws_lasalss115s2-file1.pdf
call psb_realloc(nb, acc, info)
allocate(offsets(omp_get_max_threads()))
!$omp parallel private(vals,col_inds,nnz,rwnz,thread_upperbound,acc,start_idx,end_idx) &
!$omp shared(a,b,c,offsets)
@@ -3822,66 +3946,66 @@ contains
start_idx = 0
!$omp do schedule(static) private(irw, jj, j)
do irw = 1, ma
if (start_idx == 0) then
start_idx = irw
end if
end_idx = irw
do jj = a%irp(irw), a%irp(irw + 1) - 1
j = a%ja(jj)
thread_upperbound = thread_upperbound + b%irp(j+1) - b%irp(j)
end do
if (start_idx == 0) then
start_idx = irw
end if
end_idx = irw
do jj = a%irp(irw), a%irp(irw + 1) - 1
j = a%ja(jj)
thread_upperbound = thread_upperbound + b%irp(j+1) - b%irp(j)
end do
end do
!$omp end do
call psb_realloc(thread_upperbound, vals, info)
call psb_realloc(thread_upperbound, col_inds, info)
! possible bottleneck
acc = 0
nnz = 0
!$omp do schedule(static) private(irw, jj, j, k)
do irw = 1, ma
rwnz = 0
do jj = a%irp(irw), a%irp(irw + 1) - 1
j = a%ja(jj)
do k = b%irp(j), b%irp(j + 1) - 1
if (acc(b%ja(k)) == 0) then
nnz = nnz + 1
rwnz = rwnz + 1
col_inds(nnz) = b%ja(k)
end if
acc(b%ja(k)) = acc(b%ja(k)) + a%val(jj) * b%val(k)
end do
rwnz = 0
do jj = a%irp(irw), a%irp(irw + 1) - 1
j = a%ja(jj)
do k = b%irp(j), b%irp(j + 1) - 1
if (acc(b%ja(k)) == 0) then
nnz = nnz + 1
rwnz = rwnz + 1
col_inds(nnz) = b%ja(k)
end if
acc(b%ja(k)) = acc(b%ja(k)) + a%val(jj) * b%val(k)
end do
call psb_qsort(col_inds(nnz - rwnz + 1:nnz))
do k = nnz - rwnz + 1, nnz
vals(k) = acc(col_inds(k))
acc(col_inds(k)) = 0
end do
c%irp(irw + 1) = rwnz
end do
call psb_qsort(col_inds(nnz - rwnz + 1:nnz))
do k = nnz - rwnz + 1, nnz
vals(k) = acc(col_inds(k))
acc(col_inds(k)) = 0
end do
c%irp(irw + 1) = rwnz
end do
!$omp end do
offsets(omp_get_thread_num() + 1) = nnz
!$omp barrier
! possible bottleneck
!$omp single
do k = 1, omp_get_num_threads() - 1
offsets(k + 1) = offsets(k + 1) + offsets(k)
offsets(k + 1) = offsets(k + 1) + offsets(k)
end do
!$omp end single
!$omp barrier
if (omp_get_thread_num() /= 0) then
c%irp(start_idx) = offsets(omp_get_thread_num()) + 1
c%irp(start_idx) = offsets(omp_get_thread_num()) + 1
end if
do irw = start_idx, end_idx - 1
c%irp(irw + 1) = c%irp(irw + 1) + c%irp(irw)
c%irp(irw + 1) = c%irp(irw + 1) + c%irp(irw)
end do
!$omp barrier
@@ -3891,7 +4015,7 @@ contains
call psb_realloc(c%irp(ma + 1), c%val, info)
call psb_realloc(c%irp(ma + 1), c%ja, info)
!$omp end single
c%val(c%irp(start_idx):c%irp(end_idx + 1) - 1) = vals(1:nnz)
c%ja(c%irp(start_idx):c%irp(end_idx + 1) - 1) = col_inds(1:nnz)
!$omp end parallel
@@ -3904,13 +4028,13 @@ contains
type(psb_d_csr_sparse_mat), intent(in) :: a,b
type(psb_d_csr_sparse_mat), intent(out):: c
integer(psb_ipk_), intent(out) :: info
real(psb_dpk_), allocatable :: vals(:), acc(:)
integer(psb_ipk_) :: ma, nb
integer(psb_ipk_), allocatable :: col_inds(:), offsets(:)
integer(psb_ipk_) :: irw, jj, j, k, nnz, rwnz, thread_upperbound, &
start_idx, end_idx , blk, blk_size, rwstart,&
rwblk, rwblkrem, nblks
start_idx, end_idx , blk, blk_size, rwstart,&
rwblk, rwblkrem, nblks
ma = a%get_nrows()
nb = b%get_ncols()
@@ -3931,76 +4055,76 @@ contains
start_idx = 0
!$omp do schedule(static) private(irw, jj, j)
do irw = 1, ma
do jj = a%irp(irw), a%irp(irw + 1) - 1
j = a%ja(jj)
thread_upperbound = thread_upperbound + b%irp(j+1) - b%irp(j)
end do
do jj = a%irp(irw), a%irp(irw + 1) - 1
j = a%ja(jj)
thread_upperbound = thread_upperbound + b%irp(j+1) - b%irp(j)
end do
end do
!$omp end do
call psb_realloc(thread_upperbound, vals, info)
call psb_realloc(thread_upperbound, col_inds, info)
! possible bottleneck
acc = 0
nnz = 0
!$omp do schedule(static) private(irw,jj,j,k,rwnz,blk,blk_size,rwstart)
do blk = 0, nblks - 1
if (blk < rwblkrem) then
blk_size = rwblk + 1
rwstart = blk * rwblk + blk + 1
else
blk_size = rwblk
rwstart = blk * rwblk &
+ rwblkrem + 1
if (blk < rwblkrem) then
blk_size = rwblk + 1
rwstart = blk * rwblk + blk + 1
else
blk_size = rwblk
rwstart = blk * rwblk &
+ rwblkrem + 1
end if
do irw = rwstart, rwstart + blk_size - 1
if (start_idx == 0) then
start_idx = irw
end if
do irw = rwstart, rwstart + blk_size - 1
if (start_idx == 0) then
start_idx = irw
end_idx = irw
rwnz = 0
do jj = a%irp(irw), a%irp(irw + 1) - 1
j = a%ja(jj)
do k = b%irp(j), b%irp(j + 1) - 1
if (acc(b%ja(k)) == 0) then
nnz = nnz + 1
rwnz = rwnz + 1
col_inds(nnz) = b%ja(k)
end if
end_idx = irw
rwnz = 0
do jj = a%irp(irw), a%irp(irw + 1) - 1
j = a%ja(jj)
do k = b%irp(j), b%irp(j + 1) - 1
if (acc(b%ja(k)) == 0) then
nnz = nnz + 1
rwnz = rwnz + 1
col_inds(nnz) = b%ja(k)
end if
acc(b%ja(k)) = acc(b%ja(k)) + a%val(jj) * b%val(k)
end do
end do
call psb_qsort(col_inds(nnz - rwnz + 1:nnz))
do k = nnz - rwnz + 1, nnz
vals(k) = acc(col_inds(k))
acc(col_inds(k)) = 0
end do
c%irp(irw + 1) = rwnz
acc(b%ja(k)) = acc(b%ja(k)) + a%val(jj) * b%val(k)
end do
end do
call psb_qsort(col_inds(nnz - rwnz + 1:nnz))
do k = nnz - rwnz + 1, nnz
vals(k) = acc(col_inds(k))
acc(col_inds(k)) = 0
end do
c%irp(irw + 1) = rwnz
end do
end do
!$omp end do
offsets(omp_get_thread_num() + 1) = nnz
!$omp barrier
! possible bottleneck
!$omp single
do k = 1, omp_get_num_threads() - 1
offsets(k + 1) = offsets(k + 1) + offsets(k)
offsets(k + 1) = offsets(k + 1) + offsets(k)
end do
!$omp end single
!$omp barrier
if (omp_get_thread_num() /= 0) then
c%irp(start_idx) = offsets(omp_get_thread_num()) + 1
c%irp(start_idx) = offsets(omp_get_thread_num()) + 1
end if
do irw = start_idx, end_idx - 1
c%irp(irw + 1) = c%irp(irw + 1) + c%irp(irw)
c%irp(irw + 1) = c%irp(irw + 1) + c%irp(irw)
end do
!$omp barrier
@@ -4010,7 +4134,7 @@ contains
call psb_realloc(c%irp(ma + 1), c%val, info)
call psb_realloc(c%irp(ma + 1), c%ja, info)
!$omp end single
c%val(c%irp(start_idx):c%irp(end_idx + 1) - 1) = vals(1:nnz)
c%ja(c%irp(start_idx):c%irp(end_idx + 1) - 1) = col_inds(1:nnz)
!$omp end parallel
@@ -4034,20 +4158,20 @@ contains
call c%allocate(a_m, b_n)
do row = 1, a_m
row_accs(row)%nnz = 0
nullify(row_accs(row)%root)
do col = a%irp(row), a%irp(row + 1) - 1
call psb_rb_idx_tree_scalar_sparse_row_mul(row_accs(row), a%val(col), b, a%ja(col))
end do
row_accs(row)%nnz = 0
nullify(row_accs(row)%root)
do col = a%irp(row), a%irp(row + 1) - 1
call psb_rb_idx_tree_scalar_sparse_row_mul(row_accs(row), a%val(col), b, a%ja(col))
end do
end do
call psb_rb_idx_tree_merge(row_accs, c)
deallocate(row_accs)
info = 0
end subroutine spmm_serial_rb_tree
end subroutine spmm_serial_rb_tree
subroutine spmm_omp_rb_tree(a,b,c,info)
subroutine spmm_omp_rb_tree(a,b,c,info)
use omp_lib
use psb_rb_idx_tree_mod
implicit none
@@ -4070,11 +4194,11 @@ subroutine spmm_omp_rb_tree(a,b,c,info)
!$omp parallel do schedule(static)
do row = 1, a_m
row_accs(row)%nnz = 0
nullify(row_accs(row)%root)
do col = a%irp(row), a%irp(row + 1) - 1
call psb_rb_idx_tree_scalar_sparse_row_mul(row_accs(row), a%val(col), b, a%ja(col))
end do
row_accs(row)%nnz = 0
nullify(row_accs(row)%root)
do col = a%irp(row), a%irp(row + 1) - 1
call psb_rb_idx_tree_scalar_sparse_row_mul(row_accs(row), a%val(col), b, a%ja(col))
end do
end do
!$omp end parallel do
@@ -4082,9 +4206,9 @@ subroutine spmm_omp_rb_tree(a,b,c,info)
deallocate(row_accs)
info = 0
end subroutine spmm_omp_rb_tree
end subroutine spmm_omp_rb_tree
subroutine compute_indices(a, b, c, info)
subroutine compute_indices(a, b, c, info)
implicit none
type(psb_d_csr_sparse_mat), intent(in) :: a,b
type(psb_d_csr_sparse_mat), intent(out):: c
@@ -4096,10 +4220,10 @@ subroutine compute_indices(a, b, c, info)
full_mat_bound = 0
!omp parallel do schedule(static) reduction(+:full_mat_bound)
do row = 1, a%get_nrows()
do col = a%irp(row), a%irp(row + 1) - 1
j = a%ja(col)
full_mat_bound = full_mat_bound + b%irp(j+1) - b%irp(j)
end do
do col = a%irp(row), a%irp(row + 1) - 1
j = a%ja(col)
full_mat_bound = full_mat_bound + b%irp(j+1) - b%irp(j)
end do
end do
!omp end parallel do
@@ -4109,22 +4233,22 @@ subroutine compute_indices(a, b, c, info)
c%irp(1) = 1
nnz = 0
do row = 1, a%get_nrows()
do col = a%irp(row), a%irp(row + 1) - 1
do i = b%irp(a%ja(col)), b%irp(a%ja(col) + 1) - 1
k = 0
do while(c%ja(c%irp(row) + k) /= 0 .and. c%ja(c%irp(row) + k) /= b%ja(i))
k = k + 1
end do
if (c%ja(c%irp(row) + k) == 0) then
c%ja(c%irp(row)+k) = b%ja(i)
nnz = nnz + 1
end if
end do
do col = a%irp(row), a%irp(row + 1) - 1
do i = b%irp(a%ja(col)), b%irp(a%ja(col) + 1) - 1
k = 0
do while(c%ja(c%irp(row) + k) /= 0 .and. c%ja(c%irp(row) + k) /= b%ja(i))
k = k + 1
end do
if (c%ja(c%irp(row) + k) == 0) then
c%ja(c%irp(row)+k) = b%ja(i)
nnz = nnz + 1
end if
end do
c%irp(row + 1) = nnz + 1
call psb_qsort(c%ja(c%irp(row):c%irp(row + 1)-1))
end do
c%irp(row + 1) = nnz + 1
call psb_qsort(c%ja(c%irp(row):c%irp(row + 1)-1))
end do
@@ -4132,9 +4256,9 @@ subroutine compute_indices(a, b, c, info)
call psb_realloc(nnz, c%val, info)
c%val = 0
end subroutine compute_indices
end subroutine compute_indices
subroutine direct_scalar_sparse_row_mul(out_mat, out_row_num, scalar, mat, trgt_row_num)
subroutine direct_scalar_sparse_row_mul(out_mat, out_row_num, scalar, mat, trgt_row_num)
type(psb_d_csr_sparse_mat), intent(inout) :: out_mat
integer(psb_ipk_), intent(in) :: out_row_num
real(psb_dpk_), intent(in) :: scalar
@@ -4147,17 +4271,17 @@ subroutine direct_scalar_sparse_row_mul(out_mat, out_row_num, scalar, mat, trgt_
row_end = out_mat%irp(out_row_num + 1) - 1
do i = mat%irp(trgt_row_num), mat%irp(trgt_row_num + 1) - 1
do k = out_mat%irp(out_row_num), out_mat%irp(out_row_num + 1) - 1
if (out_mat%ja(k) == mat%ja(i)) then
out_mat%val(k) = out_mat%val(k) + scalar * mat%val(i)
exit
end if
end do
do k = out_mat%irp(out_row_num), out_mat%irp(out_row_num + 1) - 1
if (out_mat%ja(k) == mat%ja(i)) then
out_mat%val(k) = out_mat%val(k) + scalar * mat%val(i)
exit
end if
end do
end do
end subroutine direct_scalar_sparse_row_mul
end subroutine direct_scalar_sparse_row_mul
subroutine spmm_omp_two_pass(a,b,c,info)
subroutine spmm_omp_two_pass(a,b,c,info)
use omp_lib
implicit none
@@ -4171,20 +4295,20 @@ subroutine spmm_omp_two_pass(a,b,c,info)
b_n = b%get_ncols()
call c%allocate(a_m, b_n)
call compute_indices(a, b, c, info)
!$omp parallel do schedule(static)
do row = 1, a_m
do col = a%irp(row), a%irp(row + 1) - 1
call direct_scalar_sparse_row_mul(c, row, a%val(col), b, a%ja(col))
end do
do col = a%irp(row), a%irp(row + 1) - 1
call direct_scalar_sparse_row_mul(c, row, a%val(col), b, a%ja(col))
end do
end do
!$omp end parallel do
end subroutine spmm_omp_two_pass
end subroutine spmm_omp_two_pass
end subroutine psb_dcsrspspmm
#endif
!
!