570 : _index_maps({p.index_map(0),
571 std::make_shared<common::IndexMap>(p.column_index_map())}),
572 _block_mode(mode), _bs({p.block_size(0), p.block_size(1)}),
573 _data(p.num_nonzeros() * _bs[0] * _bs[1], 0),
574 _cols(p.graph().first.begin(), p.graph().first.end()),
575 _row_ptr(p.graph().second.begin(), p.graph().second.end()),
581 for (
int i = 0; i < 2; ++i)
583 auto im = _index_maps[i];
584 std::int32_t size_local = im->size_local() * _bs[i];
585 std::span ghost_i = im->ghosts();
586 std::vector<std::int64_t> ghosts;
587 const std::vector<int> ghost_owner_i(im->owners().begin(),
589 std::vector<int> src_rank;
590 for (std::size_t j = 0; j < ghost_i.size(); ++j)
592 for (
int k = 0; k < _bs[i]; ++k)
594 ghosts.push_back(ghost_i[j] * _bs[i] + k);
595 src_rank.push_back(ghost_owner_i[j]);
599 std::array<std::vector<int>, 2> src_dest0
600 = {std::vector(_index_maps[i]->src().begin(),
601 _index_maps[i]->src().end()),
602 std::vector(_index_maps[i]->dest().begin(),
603 _index_maps[i]->dest().end())};
604 _index_maps[i] = std::make_shared<common::IndexMap>(
605 _index_maps[i]->comm(), size_local, src_dest0, ghosts, src_rank);
610 column_container_type new_cols;
611 new_cols.reserve(_data.size());
612 rowptr_container_type new_row_ptr{0};
613 new_row_ptr.reserve(_row_ptr.size() * _bs[0]);
614 std::span<const std::int32_t> num_diag_nnz = p.off_diagonal_offsets();
615 for (std::size_t i = 0; i < _row_ptr.size() - 1; ++i)
618 for (
int q0 = 0; q0 < _bs[0]; ++q0)
620 _off_diagonal_offset.push_back(new_row_ptr.back()
621 + num_diag_nnz[i] * _bs[1]);
622 for (
auto j = _row_ptr[i]; j < _row_ptr[i + 1]; ++j)
624 for (
int q1 = 0; q1 < _bs[1]; ++q1)
625 new_cols.push_back(_cols[j] * _bs[1] + q1);
627 new_row_ptr.push_back(new_cols.size());
631 _row_ptr = new_row_ptr;
638 std::span<const std::int32_t> num_diag_nnz = p.off_diagonal_offsets();
639 _off_diagonal_offset.reserve(num_diag_nnz.size());
640 std::ranges::transform(num_diag_nnz, _row_ptr,
641 std::back_inserter(_off_diagonal_offset),
646 std::array local_size
647 = {_index_maps[0]->size_local(), _index_maps[1]->size_local()};
649 = {_index_maps[0]->local_range(), _index_maps[1]->local_range()};
650 std::span ghosts1 = _index_maps[1]->ghosts();
652 std::span ghosts0 = _index_maps[0]->ghosts();
653 std::span src_ranks = _index_maps[0]->src();
654 std::span dest_ranks = _index_maps[0]->dest();
658 MPI_Dist_graph_create_adjacent(_index_maps[0]->comm(), dest_ranks.size(),
659 dest_ranks.data(), MPI_UNWEIGHTED,
660 src_ranks.size(), src_ranks.data(),
661 MPI_UNWEIGHTED, MPI_INFO_NULL,
false, &comm);
662 _comm = dolfinx::MPI::Comm(comm,
false);
666 _ghost_row_to_rank.reserve(_index_maps[0]->owners().
size());
667 for (
int r : _index_maps[0]->owners())
669 auto it = std::ranges::lower_bound(src_ranks, r);
670 assert(it != src_ranks.end() and *it == r);
671 std::size_t pos = std::distance(src_ranks.begin(), it);
672 _ghost_row_to_rank.push_back(pos);
676 std::vector<std::int32_t> data_per_proc(src_ranks.size(), 0);
677 for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
679 assert(_ghost_row_to_rank[i] < (
int)data_per_proc.size());
680 std::size_t pos = local_size[0] + i;
681 data_per_proc[_ghost_row_to_rank[i]] += _row_ptr[pos + 1] - _row_ptr[pos];
685 _val_send_disp.resize(src_ranks.size() + 1, 0);
686 std::partial_sum(data_per_proc.begin(), data_per_proc.end(),
687 std::next(_val_send_disp.begin()));
690 std::vector<std::int64_t> ghost_index_data(2 * _val_send_disp.back());
692 std::vector<int> insert_pos = _val_send_disp;
693 for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
695 int rank = _ghost_row_to_rank[i];
696 std::int32_t row_id = local_size[0] + i;
697 for (
int j = _row_ptr[row_id]; j < _row_ptr[row_id + 1]; ++j)
700 std::int32_t idx_pos = 2 * insert_pos[
rank];
703 ghost_index_data[idx_pos] = ghosts0[i];
704 if (std::int32_t col_local = _cols[j]; col_local < local_size[1])
705 ghost_index_data[idx_pos + 1] = col_local + local_range[1][0];
707 ghost_index_data[idx_pos + 1] = ghosts1[col_local - local_size[1]];
709 insert_pos[
rank] += 1;
715 std::vector<std::int64_t> ghost_index_array;
716 std::vector<int> recv_disp;
718 std::vector<int> send_sizes;
719 std::ranges::transform(data_per_proc, std::back_inserter(send_sizes),
720 [](
auto x) {
return 2 * x; });
722 std::vector<int> recv_sizes(dest_ranks.size());
723 send_sizes.reserve(1);
724 recv_sizes.reserve(1);
725 MPI_Neighbor_alltoall(send_sizes.data(), 1, MPI_INT, recv_sizes.data(), 1,
726 MPI_INT, _comm.comm());
729 std::vector<int> send_disp{0};
730 std::partial_sum(send_sizes.begin(), send_sizes.end(),
731 std::back_inserter(send_disp));
733 std::partial_sum(recv_sizes.begin(), recv_sizes.end(),
734 std::back_inserter(recv_disp));
736 ghost_index_array.resize(recv_disp.back());
737 MPI_Neighbor_alltoallv(ghost_index_data.data(), send_sizes.data(),
738 send_disp.data(), MPI_INT64_T,
739 ghost_index_array.data(), recv_sizes.data(),
740 recv_disp.data(), MPI_INT64_T, _comm.comm());
745 _val_recv_disp.resize(recv_disp.size());
746 int bs2 = _bs[0] * _bs[1];
747 std::ranges::transform(recv_disp, _val_recv_disp.begin(),
748 [&bs2](
auto d) { return bs2 * d / 2; });
749 std::ranges::transform(_val_send_disp, _val_send_disp.begin(),
750 [&bs2](
auto d) { return d * bs2; });
753 std::vector<std::pair<std::int64_t, std::int32_t>> global_to_local;
754 global_to_local.reserve(ghosts1.size());
755 for (std::int64_t idx : ghosts1)
756 global_to_local.push_back({idx, global_to_local.size() + local_size[1]});
757 std::ranges::sort(global_to_local);
761 for (std::size_t i = 0; i < ghost_index_array.size(); i += 2)
764 std::int32_t local_row = ghost_index_array[i] -
local_range[0][0];
765 assert(local_row >= 0 and local_row < local_size[0]);
768 std::int32_t local_col = ghost_index_array[i + 1] -
local_range[1][0];
769 if (local_col < 0 or local_col >= local_size[1])
771 auto it = std::ranges::lower_bound(
772 global_to_local, std::pair(ghost_index_array[i + 1], -1),
773 [](
auto a,
auto b) {
return a.first < b.first; });
774 assert(it != global_to_local.end()
775 and it->first == ghost_index_array[i + 1]);
776 local_col = it->second;
778 auto cit0 = std::next(_cols.begin(), _row_ptr[local_row]);
779 auto cit1 = std::next(_cols.begin(), _row_ptr[local_row + 1]);
782 auto cit = std::lower_bound(cit0, cit1, local_col);
784 assert(*cit == local_col);
785 std::size_t d = std::distance(_cols.begin(), cit);
786 _unpack_pos.push_back(d);
789 _unpack_pos.shrink_to_fit();