impure function mat_vec_coarray_rel(A, v, transA, option, coarray) result(w)
real(rk), intent(in), contiguous :: A(:,:), v(:)
character(*), intent(in), optional :: option
logical, intent(in), optional :: transA
real(rk), allocatable :: w(:)
logical, intent(in) :: coarray
#if defined (USE_COARRAY)
integer :: i, im, nimg, n, m
integer :: block_size(num_images()), start_elem(num_images()), end_elem(num_images())
real(rk), allocatable :: w_block(:)[:], v_block(:)[:], A_block(:,:)[:]
if (present(transA)) then
if (transA) then
! ATv
allocate(w(size(A,2)), source=0.0_rk)
im = this_image()
nimg = num_images()
call compute_block_ranges(size(A,2), nimg, block_size, start_elem, end_elem)
allocate(w_block(block_size(im))[*], A_block(size(A,1), block_size(im))[*])
A_block(:,:)[im] = A(:, start_elem(im):end_elem(im))
w_block(:)[im] = matmul(A_block(:, :)[im], v, transA=.true., option=option)
sync all
if (im == 1) then
do i = 1, nimg
w(start_elem(i):end_elem(i)) = w_block(:)[i]
end do
end if
else if (.not. transA) then
! Av
allocate(w(size(A,1)), source=0.0_rk)
im = this_image()
nimg = num_images()
call compute_block_ranges(size(A,2), nimg, block_size, start_elem, end_elem)
allocate(w_block(size(A,1))[*], v_block(block_size(im))[*], A_block(size(A,1), block_size(im))[*])
A_block(:,:)[im] = A(:, start_elem(im):end_elem(im))
v_block(:)[im] = v(start_elem(im):end_elem(im))
w_block(:)[im] = matmul(A_block(:,:)[im], v_block(:)[im], transA=.false., option=option)
sync all
if (im == 1) then
do i = 1, nimg
w(:) = w(:) + w_block(:)[i]
end do
end if
end if
else if (.not. present(transA)) then
! Av
allocate(w(size(A,1)), source=0.0_rk)
im = this_image()
nimg = num_images()
call compute_block_ranges(size(A,2), nimg, block_size, start_elem, end_elem)
allocate(w_block(size(A,1))[*], v_block(block_size(im))[*], A_block(size(A,1), block_size(im))[*])
A_block(:,:)[im] = A(:, start_elem(im):end_elem(im))
v_block(:)[im] = v(start_elem(im):end_elem(im))
w_block(:)[im] = matmul(A_block(:,:)[im], v_block(:)[im], transA=.false., option=option)
sync all
if (im == 1) then
do i = 1, nimg
w(:) = w(:) + w_block(:)[i]
end do
end if
end if
#else
w = matmul(A, v, transA=transA, option=option)
#endif
end function mat_vec_coarray_rel