522 : _index_maps({p.index_map(0),
523 std::make_shared<common::IndexMap>(p.column_index_map())}),
524 _block_mode(mode), _bs({p.block_size(0), p.block_size(1)}),
525 _data(p.num_nonzeros() * _bs[0] * _bs[1], 0),
526 _cols(p.graph().first.begin(), p.graph().first.end()),
527 _row_ptr(p.graph().second.begin(), p.graph().second.end()),
533 for (
int i = 0; i < 2; ++i)
535 const auto im = _index_maps[i];
536 const std::int32_t size_local = im->size_local() * _bs[i];
537 std::span ghost_i = im->ghosts();
538 std::vector<std::int64_t> ghosts;
539 const std::vector<int> ghost_owner_i(im->owners().begin(),
541 std::vector<int> src_rank;
542 for (std::size_t j = 0; j < ghost_i.size(); ++j)
544 for (
int k = 0; k < _bs[i]; ++k)
546 ghosts.push_back(ghost_i[j] * _bs[i] + k);
547 src_rank.push_back(ghost_owner_i[j]);
551 std::array<std::vector<int>, 2> src_dest0
552 = {std::vector(_index_maps[i]->src().begin(),
553 _index_maps[i]->src().end()),
554 std::vector(_index_maps[i]->dest().begin(),
555 _index_maps[i]->dest().end())};
556 _index_maps[i] = std::make_shared<common::IndexMap>(
557 _index_maps[i]->comm(), size_local, src_dest0, ghosts, src_rank);
562 column_container_type new_cols;
563 new_cols.reserve(_data.size());
564 rowptr_container_type new_row_ptr = {0};
565 new_row_ptr.reserve(_row_ptr.size() * _bs[0]);
566 std::span<const std::int32_t> num_diag_nnz = p.off_diagonal_offsets();
567 for (std::size_t i = 0; i < _row_ptr.size() - 1; ++i)
570 for (
int q0 = 0; q0 < _bs[0]; ++q0)
572 _off_diagonal_offset.push_back(new_row_ptr.back()
573 + num_diag_nnz[i] * _bs[1]);
574 for (
auto j = _row_ptr[i]; j < _row_ptr[i + 1]; ++j)
576 for (
int q1 = 0; q1 < _bs[1]; ++q1)
577 new_cols.push_back(_cols[j] * _bs[1] + q1);
579 new_row_ptr.push_back(new_cols.size());
583 _row_ptr = new_row_ptr;
590 std::span<const std::int32_t> num_diag_nnz = p.off_diagonal_offsets();
591 _off_diagonal_offset.reserve(num_diag_nnz.size());
592 std::ranges::transform(num_diag_nnz, _row_ptr,
593 std::back_inserter(_off_diagonal_offset),
598 const std::array local_size
599 = {_index_maps[0]->size_local(), _index_maps[1]->size_local()};
601 = {_index_maps[0]->local_range(), _index_maps[1]->local_range()};
602 std::span ghosts1 = _index_maps[1]->ghosts();
604 std::span ghosts0 = _index_maps[0]->ghosts();
605 std::span src_ranks = _index_maps[0]->src();
606 std::span dest_ranks = _index_maps[0]->dest();
610 MPI_Dist_graph_create_adjacent(_index_maps[0]->comm(), dest_ranks.size(),
611 dest_ranks.data(), MPI_UNWEIGHTED,
612 src_ranks.size(), src_ranks.data(),
613 MPI_UNWEIGHTED, MPI_INFO_NULL,
false, &comm);
614 _comm = dolfinx::MPI::Comm(comm,
false);
618 _ghost_row_to_rank.reserve(_index_maps[0]->owners().
size());
619 for (
int r : _index_maps[0]->owners())
621 auto it = std::ranges::lower_bound(src_ranks, r);
622 assert(it != src_ranks.end() and *it == r);
623 std::size_t pos = std::distance(src_ranks.begin(), it);
624 _ghost_row_to_rank.push_back(pos);
628 std::vector<std::int32_t> data_per_proc(src_ranks.size(), 0);
629 for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
631 assert(_ghost_row_to_rank[i] < (
int)data_per_proc.size());
632 std::size_t pos = local_size[0] + i;
633 data_per_proc[_ghost_row_to_rank[i]] += _row_ptr[pos + 1] - _row_ptr[pos];
637 _val_send_disp.resize(src_ranks.size() + 1, 0);
638 std::partial_sum(data_per_proc.begin(), data_per_proc.end(),
639 std::next(_val_send_disp.begin()));
642 std::vector<std::int64_t> ghost_index_data(2 * _val_send_disp.back());
644 std::vector<int> insert_pos = _val_send_disp;
645 for (std::size_t i = 0; i < _ghost_row_to_rank.size(); ++i)
647 const int rank = _ghost_row_to_rank[i];
648 std::int32_t row_id = local_size[0] + i;
649 for (
int j = _row_ptr[row_id]; j < _row_ptr[row_id + 1]; ++j)
652 const std::int32_t idx_pos = 2 * insert_pos[
rank];
655 ghost_index_data[idx_pos] = ghosts0[i];
656 if (std::int32_t col_local = _cols[j]; col_local < local_size[1])
657 ghost_index_data[idx_pos + 1] = col_local + local_range[1][0];
659 ghost_index_data[idx_pos + 1] = ghosts1[col_local - local_size[1]];
661 insert_pos[
rank] += 1;
667 std::vector<std::int64_t> ghost_index_array;
668 std::vector<int> recv_disp;
670 std::vector<int> send_sizes;
671 std::ranges::transform(data_per_proc, std::back_inserter(send_sizes),
672 [](
auto x) {
return 2 * x; });
674 std::vector<int> recv_sizes(dest_ranks.size());
675 send_sizes.reserve(1);
676 recv_sizes.reserve(1);
677 MPI_Neighbor_alltoall(send_sizes.data(), 1, MPI_INT, recv_sizes.data(), 1,
678 MPI_INT, _comm.comm());
681 std::vector<int> send_disp = {0};
682 std::partial_sum(send_sizes.begin(), send_sizes.end(),
683 std::back_inserter(send_disp));
685 std::partial_sum(recv_sizes.begin(), recv_sizes.end(),
686 std::back_inserter(recv_disp));
688 ghost_index_array.resize(recv_disp.back());
689 MPI_Neighbor_alltoallv(ghost_index_data.data(), send_sizes.data(),
690 send_disp.data(), MPI_INT64_T,
691 ghost_index_array.data(), recv_sizes.data(),
692 recv_disp.data(), MPI_INT64_T, _comm.comm());
697 _val_recv_disp.resize(recv_disp.size());
698 const int bs2 = _bs[0] * _bs[1];
699 std::ranges::transform(recv_disp, _val_recv_disp.begin(),
700 [&bs2](
auto d) { return bs2 * d / 2; });
701 std::ranges::transform(_val_send_disp, _val_send_disp.begin(),
702 [&bs2](
auto d) { return d * bs2; });
705 std::vector<std::pair<std::int64_t, std::int32_t>> global_to_local;
706 global_to_local.reserve(ghosts1.size());
707 for (std::int64_t idx : ghosts1)
708 global_to_local.push_back({idx, global_to_local.size() + local_size[1]});
709 std::ranges::sort(global_to_local);
713 for (std::size_t i = 0; i < ghost_index_array.size(); i += 2)
716 const std::int32_t local_row = ghost_index_array[i] -
local_range[0][0];
717 assert(local_row >= 0 and local_row < local_size[0]);
720 std::int32_t local_col = ghost_index_array[i + 1] -
local_range[1][0];
721 if (local_col < 0 or local_col >= local_size[1])
723 auto it = std::ranges::lower_bound(
724 global_to_local, std::pair(ghost_index_array[i + 1], -1),
725 [](
auto& a,
auto& b) {
return a.first < b.first; });
726 assert(it != global_to_local.end()
727 and it->first == ghost_index_array[i + 1]);
728 local_col = it->second;
730 auto cit0 = std::next(_cols.begin(), _row_ptr[local_row]);
731 auto cit1 = std::next(_cols.begin(), _row_ptr[local_row + 1]);
734 auto cit = std::lower_bound(cit0, cit1, local_col);
736 assert(*cit == local_col);
737 std::size_t d = std::distance(_cols.begin(), cit);
738 _unpack_pos.push_back(d);
775 std::span<const std::int64_t> Arow_ptr(
row_ptr().data(), nrowslocal + 1);
776 std::span<const std::int32_t> Acols(
cols().data(), Arow_ptr[nrowslocal]);
777 std::span<const std::int64_t> Aoff_diag_offset(
off_diag_offset().data(),
779 std::span<const Scalar> Avalues(
values().data(), Arow_ptr[nrowslocal]);
781 std::span<const Scalar> _x = x.
array();
784 std::span<const std::int64_t> Arow_begin(Arow_ptr.data(), nrowslocal);
785 std::span<const std::int64_t> Arow_end(Arow_ptr.data() + 1, nrowslocal);
791 impl::spmv<Scalar, 1>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x, _y,
796 impl::spmv<Scalar, -1>(Avalues, Arow_begin, Aoff_diag_offset, Acols, _x, _y,
807 impl::spmv<Scalar, 1>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x, _y,
812 impl::spmv<Scalar, -1>(Avalues, Aoff_diag_offset, Arow_end, Acols, _x, _y,