563 : _index_maps({p.index_map(0),
564 std::make_shared<common::IndexMap>(p.column_index_map())}),
565 _block_mode(mode), _bs({p.block_size(0), p.block_size(1)}),
566 _data(p.num_nonzeros() * _bs[0] * _bs[1], 0),
567 _cols(p.graph().first.begin(), p.graph().first.end()),
568 _row_ptr(p.graph().second.begin(), p.graph().second.end()),
574 for (
int i = 0; i < 2; ++i)
576 const auto im = _index_maps[i];
577 const std::int32_t size_local = im->size_local() * _bs[i];
578 std::span ghost_i = im->ghosts();
579 std::vector<std::int64_t> ghosts;
580 const std::vector<int> ghost_owner_i(im->owners().begin(),
582 std::vector<int> src_rank;
583 for (std::size_t j = 0; j < ghost_i.size(); ++j)
585 for (
int k = 0; k < _bs[i]; ++k)
587 ghosts.push_back(ghost_i[j] * _bs[i] + k);
588 src_rank.push_back(ghost_owner_i[j]);
592 std::array<std::vector<int>, 2> src_dest0
593 = {std::vector(_index_maps[i]->src().begin(),
594 _index_maps[i]->src().end()),
595 std::vector(_index_maps[i]->dest().begin(),
596 _index_maps[i]->dest().end())};
597 _index_maps[i] = std::make_shared<common::IndexMap>(
598 _index_maps[i]->comm(), size_local, src_dest0, ghosts, src_rank);
603 column_container_type new_cols;
604 new_cols.reserve(_data.size());
605 rowptr_container_type new_row_ptr{0};
606 new_row_ptr.reserve(_row_ptr.size() * _bs[0]);
607 std::span<const std::int32_t> num_diag_nnz = p.off_diagonal_offsets();
608 for (std::size_t i = 0; i < _row_ptr.size() - 1; ++i)
611 for (
int q0 = 0; q0 < _bs[0]; ++q0)
613 _off_diagonal_offset.push_back(new_row_ptr.back()
614 + num_diag_nnz[i] * _bs[1]);
615 for (
auto j = _row_ptr[i]; j < _row_ptr[i + 1]; ++j)
617 for (
int q1 = 0; q1 < _bs[1]; ++q1)
618 new_cols.push_back(_cols[j] * _bs[1] + q1);
620 new_row_ptr.push_back(new_cols.size());
624 _row_ptr = new_row_ptr;
631 std::span<const std::int32_t> num_diag_nnz = p.off_diagonal_offsets();
632 _off_diagonal_offset.reserve(num_diag_nnz.size());
633 std::ranges::transform(num_diag_nnz, _row_ptr,
634 std::back_inserter(_off_diagonal_offset),
639 const std::array local_size
640 = {_index_maps[0]->size_local(), _index_maps[1]->size_local()};
642 = {_index_maps[0]->local_range(), _index_maps[1]->local_range()};
643 std::span ghosts1 = _index_maps[1]->ghosts();
645 std::span ghosts0 = _index_maps[0]->ghosts();
646 std::span src_ranks = _index_maps[0]->src();
647 std::span dest_ranks = _index_maps[0]->dest();
651 MPI_Dist_graph_create_adjacent(_index_maps[0]->comm(), dest_ranks.size(),
652 dest_ranks.data(), MPI_UNWEIGHTED,
653 src_ranks.size(), src_ranks.data(),
654 MPI_UNWEIGHTED, MPI_INFO_NULL,
false, &comm);
655 _comm = dolfinx::MPI::Comm(comm,
false);
659 _ghost_row_to_rank.reserve(_index_maps[0]->owners().
size());
660 for (
int r : _index_maps[0]->owners())
662 auto it = std::ranges::lower_bound(src_ranks, r);
663 assert(it != src_ranks.end() and *it == r);
664 std::size_t pos = std::distance(src_ranks.begin(), it);
665 _ghost_row_to_rank.push_back(pos);
669 std::vector<std::int32_t> data_per_proc(src_ranks.size(), 0);
670 for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
672 assert(_ghost_row_to_rank[i] < (
int)data_per_proc.size());
673 std::size_t pos = local_size[0] + i;
674 data_per_proc[_ghost_row_to_rank[i]] += _row_ptr[pos + 1] - _row_ptr[pos];
678 _val_send_disp.resize(src_ranks.size() + 1, 0);
679 std::partial_sum(data_per_proc.begin(), data_per_proc.end(),
680 std::next(_val_send_disp.begin()));
683 std::vector<std::int64_t> ghost_index_data(2 * _val_send_disp.back());
685 std::vector<int> insert_pos = _val_send_disp;
686 for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
688 const int rank = _ghost_row_to_rank[i];
689 std::int32_t row_id = local_size[0] + i;
690 for (
int j = _row_ptr[row_id]; j < _row_ptr[row_id + 1]; ++j)
693 const std::int32_t idx_pos = 2 * insert_pos[
rank];
696 ghost_index_data[idx_pos] = ghosts0[i];
697 if (std::int32_t col_local = _cols[j]; col_local < local_size[1])
698 ghost_index_data[idx_pos + 1] = col_local + local_range[1][0];
700 ghost_index_data[idx_pos + 1] = ghosts1[col_local - local_size[1]];
702 insert_pos[
rank] += 1;
708 std::vector<std::int64_t> ghost_index_array;
709 std::vector<int> recv_disp;
711 std::vector<int> send_sizes;
712 std::ranges::transform(data_per_proc, std::back_inserter(send_sizes),
713 [](
auto x) {
return 2 * x; });
715 std::vector<int> recv_sizes(dest_ranks.size());
716 send_sizes.reserve(1);
717 recv_sizes.reserve(1);
718 MPI_Neighbor_alltoall(send_sizes.data(), 1, MPI_INT, recv_sizes.data(), 1,
719 MPI_INT, _comm.comm());
722 std::vector<int> send_disp{0};
723 std::partial_sum(send_sizes.begin(), send_sizes.end(),
724 std::back_inserter(send_disp));
726 std::partial_sum(recv_sizes.begin(), recv_sizes.end(),
727 std::back_inserter(recv_disp));
729 ghost_index_array.resize(recv_disp.back());
730 MPI_Neighbor_alltoallv(ghost_index_data.data(), send_sizes.data(),
731 send_disp.data(), MPI_INT64_T,
732 ghost_index_array.data(), recv_sizes.data(),
733 recv_disp.data(), MPI_INT64_T, _comm.comm());
738 _val_recv_disp.resize(recv_disp.size());
739 const int bs2 = _bs[0] * _bs[1];
740 std::ranges::transform(recv_disp, _val_recv_disp.begin(),
741 [&bs2](
auto d) { return bs2 * d / 2; });
742 std::ranges::transform(_val_send_disp, _val_send_disp.begin(),
743 [&bs2](
auto d) { return d * bs2; });
746 std::vector<std::pair<std::int64_t, std::int32_t>> global_to_local;
747 global_to_local.reserve(ghosts1.size());
748 for (std::int64_t idx : ghosts1)
749 global_to_local.push_back({idx, global_to_local.size() + local_size[1]});
750 std::ranges::sort(global_to_local);
754 for (std::size_t i = 0; i < ghost_index_array.size(); i += 2)
757 const std::int32_t local_row = ghost_index_array[i] -
local_range[0][0];
758 assert(local_row >= 0 and local_row < local_size[0]);
761 std::int32_t local_col = ghost_index_array[i + 1] -
local_range[1][0];
762 if (local_col < 0 or local_col >= local_size[1])
764 auto it = std::ranges::lower_bound(
765 global_to_local, std::pair(ghost_index_array[i + 1], -1),
766 [](
auto& a,
auto& b) {
return a.first < b.first; });
767 assert(it != global_to_local.end()
768 and it->first == ghost_index_array[i + 1]);
769 local_col = it->second;
771 auto cit0 = std::next(_cols.begin(), _row_ptr[local_row]);
772 auto cit1 = std::next(_cols.begin(), _row_ptr[local_row + 1]);
775 auto cit = std::lower_bound(cit0, cit1, local_col);
777 assert(*cit == local_col);
778 std::size_t d = std::distance(_cols.begin(), cit);
779 _unpack_pos.push_back(d);
816 std::span<const std::int64_t> Arow_ptr(
row_ptr().data(), nrowslocal + 1);
817 std::span<const std::int32_t> Acols(
cols().data(), Arow_ptr[nrowslocal]);
818 std::span<const std::int64_t> Aoff_diag_offset(
off_diag_offset().data(),
820 std::span<const Scalar> Avalues(
values().data(), Arow_ptr[nrowslocal]);
822 std::span<const Scalar> _x = x.
array();
823 std::span<Scalar> _y = y.
array();
825 std::span<const std::int64_t> Arow_begin(Arow_ptr.data(), nrowslocal);
826 std::span<const std::int64_t> Arow_end(Arow_ptr.data() + 1, nrowslocal);
832 impl::spmv<Scalar, 1>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x, _y,
837 impl::spmv<Scalar, -1>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x, _y,
848 impl::spmv<Scalar, 1>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x, _y,
853 impl::spmv<Scalar, -1>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x, _y,