126 const int bs_rt = (BS < 0) ? B.
block_size()[0] : BS;
127 const int bsq = bs_rt * bs_rt;
130 std::span<const int> src = col_map_A->src();
131 std::span<const int> dest = col_map_A->dest();
132 MPI_Comm comm = col_map_A->comm();
145 auto new_col_map = std::make_shared<dolfinx::common::IndexMap>(
146 col_map_B->comm(), col_map_B->size_local(), col_map_B->ghosts(),
147 col_map_B->owners());
148 return {new_col_map, std::vector<std::int32_t>{0},
149 std::vector<std::int32_t>{}, std::vector<T>{}};
154 auto rank_to_nbr = [&src](
int r) ->
int
156 auto it = std::lower_bound(src.begin(), src.end(), r);
157 assert(it != src.end() && *it == r);
158 return static_cast<int>(std::ranges::distance(src.begin(), it));
160 MPI_Comm neigh_comm_fwd;
161 MPI_Comm neigh_comm_rev;
162 MPI_Dist_graph_create_adjacent(comm, dest.size(), dest.data(), MPI_UNWEIGHTED,
163 src.size(), src.data(), MPI_UNWEIGHTED,
164 MPI_INFO_NULL,
false, &neigh_comm_fwd);
165 MPI_Dist_graph_create_adjacent(comm, src.size(), src.data(), MPI_UNWEIGHTED,
166 dest.size(), dest.data(), MPI_UNWEIGHTED,
167 MPI_INFO_NULL,
false, &neigh_comm_rev);
170 std::span<const std::int64_t> required_globals_map = col_map_A->ghosts();
173 std::span<const int> ghost_owners = col_map_A->owners();
175 std::vector<int> send_count(src.size(), 0);
176 std::vector<int> recv_count(dest.size(), 0);
177 send_count.reserve(1);
178 recv_count.reserve(1);
180 for (
int gh : ghost_owners)
181 ++send_count[rank_to_nbr(gh)];
182 MPI_Neighbor_alltoall(send_count.data(), 1, MPI_INT, recv_count.data(), 1,
183 MPI_INT, neigh_comm_fwd);
186 std::vector<int> send_disp(src.size() + 1, 0);
187 std::partial_sum(send_count.begin(), send_count.end(),
188 std::next(send_disp.begin()));
189 std::vector<int> recv_disp(dest.size() + 1, 0);
190 std::partial_sum(recv_count.begin(), recv_count.end(),
191 std::next(recv_disp.begin()));
195 std::vector<std::int32_t> perm(ghost_owners.size());
196 std::vector<std::int64_t> required_globals(required_globals_map.size());
197 for (std::size_t i = 0; i < ghost_owners.size(); ++i)
199 int pos = rank_to_nbr(ghost_owners[i]);
200 perm[i] = send_disp[pos];
201 required_globals[send_disp[pos]++] = required_globals_map[i];
206 std::partial_sum(send_count.begin(), send_count.end(),
207 std::next(send_disp.begin()));
211 std::vector<std::int64_t> recv_row_globals(recv_disp.back());
213 MPI_Neighbor_alltoallv(required_globals.data(), send_count.data(),
214 send_disp.data(), MPI_INT64_T, recv_row_globals.data(),
215 recv_count.data(), recv_disp.data(), MPI_INT64_T,
221 std::vector<std::int32_t> recv_row_locals(recv_row_globals.size());
222 row_map_B->global_to_local(recv_row_globals, recv_row_locals);
227 std::vector<int> send_entry_count(dest.size(), 0);
228 std::vector<int> send_entry_disp(dest.size() + 1, 0);
229 std::vector<int> recv_entry_count(src.size(), 0);
230 std::vector<int> recv_entry_disp(src.size() + 1, 0);
231 send_entry_count.reserve(1);
232 recv_entry_count.reserve(1);
233 assert(recv_count.size() == dest.size());
234 for (std::size_t p = 0; p < recv_count.size(); ++p)
236 for (std::int32_t i = recv_disp[p]; i < recv_disp[p + 1]; ++i)
238 int r = recv_row_locals[i];
239 send_entry_count[p] += row_ptr_B[r + 1] - row_ptr_B[r];
242 std::partial_sum(send_entry_count.begin(), send_entry_count.end(),
243 std::next(send_entry_disp.begin()));
245 MPI_Neighbor_alltoall(send_entry_count.data(), 1, MPI_INT,
246 recv_entry_count.data(), 1, MPI_INT, neigh_comm_rev);
248 std::partial_sum(recv_entry_count.begin(), recv_entry_count.end(),
249 std::next(recv_entry_disp.begin()));
254 auto values_B = B.
values();
255 std::vector<T> send_vals(send_entry_disp.back() * bsq);
256 std::vector<T> recv_vals(recv_entry_disp.back() * bsq);
258 for (
int r : recv_row_locals)
260 std::size_t row_size = row_ptr_B[r + 1] - row_ptr_B[r];
261 std::copy(std::next(values_B.begin(), row_ptr_B[r] * bsq),
262 std::next(values_B.begin(), row_ptr_B[r + 1] * bsq),
263 std::next(send_vals.begin(), k));
270 MPI_Datatype mpi_block = block_type.
type();
271 MPI_Neighbor_alltoallv(send_vals.data(), send_entry_count.data(),
272 send_entry_disp.data(), mpi_block, recv_vals.data(),
273 recv_entry_count.data(), recv_entry_disp.data(),
274 mpi_block, neigh_comm_rev);
277 auto cols_B = B.
cols();
279 const std::int32_t num_owned_cols_B = col_map_B->size_local();
280 std::span<const int> ghost_owners_B = col_map_B->owners();
283 std::vector<std::int64_t> send_col_data(send_entry_disp.back());
284 std::vector<std::int64_t> recv_col_data(recv_entry_disp.back());
285 std::vector<int> send_col_owners(send_entry_disp.back());
286 std::vector<int> recv_col_owners(recv_entry_disp.back());
289 std::vector<int> send_row_size(recv_disp.back());
290 std::vector<int> recv_row_size(send_disp.back());
294 for (std::size_t p = 0; p < dest.size(); ++p)
296 for (std::int32_t i = recv_disp[p]; i < recv_disp[p + 1]; ++i)
298 int r = recv_row_locals[i];
299 std::size_t row_size = row_ptr_B[r + 1] - row_ptr_B[r];
300 col_map_B->local_to_global(
301 std::span(std::next(cols_B.begin(), row_ptr_B[r]), row_size),
302 std::span(std::next(send_col_data.begin(), k), row_size));
303 for (std::size_t j = 0; j < row_size; ++j)
305 std::int32_t local_col = cols_B[row_ptr_B[r] + j];
306 send_col_owners[k + j]
307 = (local_col < num_owned_cols_B)
309 : ghost_owners_B[local_col - num_owned_cols_B];
312 send_row_size[i] = row_size;
316 MPI_Neighbor_alltoallv(send_col_data.data(), send_entry_count.data(),
317 send_entry_disp.data(), MPI_INT64_T,
318 recv_col_data.data(), recv_entry_count.data(),
319 recv_entry_disp.data(), MPI_INT64_T, neigh_comm_rev);
320 MPI_Neighbor_alltoallv(send_col_owners.data(), send_entry_count.data(),
321 send_entry_disp.data(), MPI_INT,
322 recv_col_owners.data(), recv_entry_count.data(),
323 recv_entry_disp.data(), MPI_INT, neigh_comm_rev);
324 MPI_Neighbor_alltoallv(send_row_size.data(), recv_count.data(),
325 recv_disp.data(), MPI_INT, recv_row_size.data(),
326 send_count.data(), send_disp.data(), MPI_INT,
329 MPI_Comm_free(&neigh_comm_fwd);
330 MPI_Comm_free(&neigh_comm_rev);
335 std::vector<std::int32_t> ghost_row_ptr(recv_row_size.size() + 1, 0);
336 std::partial_sum(recv_row_size.begin(), recv_row_size.end(),
337 std::next(ghost_row_ptr.begin()));
345 auto [col_B_local_start, col_B_local_end] = col_map_B->local_range();
346 auto is_local = [col_B_local_start, col_B_local_end](std::int64_t idx)
347 {
return idx >= col_B_local_start && idx < col_B_local_end; };
350 std::vector<std::pair<std::int64_t, int>> col_owner_pairs;
351 col_owner_pairs.reserve(recv_col_data.size() + col_map_B->ghosts().size());
352 for (std::size_t i = 0; i < recv_col_data.size(); ++i)
353 if (!is_local(recv_col_data[i]))
354 col_owner_pairs.push_back({recv_col_data[i], recv_col_owners[i]});
356 std::span<const std::int64_t> existing_ghosts = col_map_B->ghosts();
357 std::span<const int> existing_ghost_owners = col_map_B->owners();
358 for (std::size_t i = 0; i < existing_ghosts.size(); ++i)
359 col_owner_pairs.push_back({existing_ghosts[i], existing_ghost_owners[i]});
362 std::sort(col_owner_pairs.begin(), col_owner_pairs.end());
363 col_owner_pairs.erase(
364 std::unique(col_owner_pairs.begin(), col_owner_pairs.end()),
365 col_owner_pairs.end());
367 std::vector<std::int64_t> unique_cols(col_owner_pairs.size());
368 std::vector<int> unique_col_owners(col_owner_pairs.size());
369 for (std::size_t i = 0; i < col_owner_pairs.size(); ++i)
371 unique_cols[i] = col_owner_pairs[i].first;
372 unique_col_owners[i] = col_owner_pairs[i].second;
376 auto new_col_map = std::make_shared<dolfinx::common::IndexMap>(
377 comm, col_map_B->size_local(), unique_cols, unique_col_owners);
380 std::vector<std::int32_t> recv_col_local(recv_col_data.size());
381 new_col_map->global_to_local(recv_col_data, recv_col_local);
386 const std::size_t num_ghosts = perm.size();
387 std::vector<std::int32_t> orig_ghost_row_ptr(num_ghosts + 1, 0);
388 for (std::size_t i = 0; i < num_ghosts; ++i)
389 orig_ghost_row_ptr[i + 1] = recv_row_size[perm[i]];
390 std::partial_sum(orig_ghost_row_ptr.begin(), orig_ghost_row_ptr.end(),
391 orig_ghost_row_ptr.begin());
393 std::vector<std::int32_t> orig_recv_col_local(recv_col_local.size());
394 std::vector<T> orig_recv_vals(recv_vals.size());
395 for (std::size_t i = 0; i < num_ghosts; ++i)
397 std::int32_t src_start = ghost_row_ptr[perm[i]];
398 std::int32_t src_end = ghost_row_ptr[perm[i] + 1];
399 std::int32_t dst_start = orig_ghost_row_ptr[i];
400 std::copy(std::next(recv_col_local.begin(), src_start),
401 std::next(recv_col_local.begin(), src_end),
402 std::next(orig_recv_col_local.begin(), dst_start));
403 std::copy(std::next(recv_vals.begin(), src_start * bsq),
404 std::next(recv_vals.begin(), src_end * bsq),
405 std::next(orig_recv_vals.begin(), dst_start * bsq));
408 return {new_col_map, orig_ghost_row_ptr, orig_recv_col_local, orig_recv_vals};
436 std::shared_ptr<const common::IndexMap> new_col_map,
437 std::span<const std::int32_t> ghost_row_ptr,
438 std::span<const std::int32_t> ghost_cols, std::span<const T> ghost_vals)
441 const int bs_rt = (BS < 0) ? A.
block_size()[0] : BS;
442 const int bsq = bs_rt * bs_rt;
445 const std::int32_t num_local_rows_A = A.
index_map(0)->size_local();
446 const std::int32_t num_local_rows_B = B.
index_map(0)->size_local();
447 const std::int32_t num_owned_cols_B = col_map_B->size_local();
448 const std::int32_t num_owned_cols_C = new_col_map->size_local();
452 auto cols_A = A.
cols();
455 auto cols_B = B.
cols();
460 std::span<const std::int64_t> b_ghost_globals = col_map_B->ghosts();
461 std::vector<std::int32_t> b_ghost_remap(b_ghost_globals.size());
462 new_col_map->global_to_local(b_ghost_globals, b_ghost_remap);
464 const std::int32_t num_cols_C
465 = new_col_map->size_local()
466 +
static_cast<std::int32_t
>(new_col_map->ghosts().size());
468 std::vector<std::int64_t> row_ptr_C = {0};
469 row_ptr_C.reserve(num_local_rows_A + 1);
470 std::vector<std::int32_t> off_diag_offsets_C;
471 off_diag_offsets_C.reserve(num_local_rows_A);
472 std::vector<std::int32_t> cols_C;
473 std::vector<T> vals_C;
479 std::vector<T> acc(num_cols_C * bsq, T(0));
480 std::vector<bool> in_row(num_cols_C,
false);
481 std::vector<std::int32_t> row_cols;
486 auto gemm_block = [&](
const T* a_blk,
const T* b_blk, T* c_blk)
488 if constexpr (BS < 0)
490 for (
int r = 0; r < bs_rt; ++r)
491 for (
int t = 0; t < bs_rt; ++t)
493 const T a_rt = a_blk[r * bs_rt + t];
494 for (
int s = 0; s < bs_rt; ++s)
495 c_blk[r * bs_rt + s] += a_rt * b_blk[t * bs_rt + s];
500 for (
int r = 0; r < BS; ++r)
501 for (
int t = 0; t < BS; ++t)
503 const T a_rt = a_blk[r * BS + t];
504 for (
int s = 0; s < BS; ++s)
505 c_blk[r * BS + s] += a_rt * b_blk[t * BS + s];
510 for (std::int32_t i = 0; i < num_local_rows_A; ++i)
515 for (std::int32_t ka = row_ptr_A[i]; ka < offdiag_ptr_A[i]; ++ka)
517 const std::int32_t j = cols_A[ka];
518 const T* a_blk = vals_A.data() + ka * bsq;
519 for (std::int32_t kb = row_ptr_B[j]; kb < row_ptr_B[j + 1]; ++kb)
521 const std::int32_t c = cols_B[kb];
523 = c < num_owned_cols_B ? c : b_ghost_remap[c - num_owned_cols_B];
527 row_cols.push_back(k);
529 gemm_block(a_blk, vals_B.data() + kb * bsq, acc.data() + k * bsq);
534 for (std::int32_t ka = offdiag_ptr_A[i]; ka < row_ptr_A[i + 1]; ++ka)
536 const std::int32_t j = cols_A[ka];
537 const T* a_blk = vals_A.data() + ka * bsq;
538 const std::int32_t g = j - num_local_rows_B;
539 for (std::int32_t kb = ghost_row_ptr[g]; kb < ghost_row_ptr[g + 1]; ++kb)
541 const std::int32_t k = ghost_cols[kb];
545 row_cols.push_back(k);
547 gemm_block(a_blk, ghost_vals.data() + kb * bsq, acc.data() + k * bsq);
553 auto pos = row_cols.begin();
554 while (pos != row_cols.end())
556 const T* c_blk = acc.data() + (*pos) * bsq;
557 bool all_zero =
true;
558 for (
int q = 0; q < bsq && all_zero; ++q)
559 all_zero = (c_blk[q] == T(0));
562 in_row[*pos] =
false;
563 pos = row_cols.erase(pos);
571 std::sort(row_cols.begin(), row_cols.end());
575 off_diag_offsets_C.push_back(
static_cast<std::int32_t
>(
576 std::lower_bound(row_cols.begin(), row_cols.end(), num_owned_cols_C)
577 - row_cols.begin()));
580 for (
const std::int32_t c : row_cols)
583 T* c_blk = acc.data() + c * bsq;
584 for (
int q = 0; q < bsq; ++q)
586 vals_C.push_back(c_blk[q]);
591 row_ptr_C.push_back(
static_cast<std::int64_t
>(cols_C.size()));
594 return {std::move(row_ptr_C), std::move(off_diag_offsets_C),
595 std::move(cols_C), std::move(vals_C)};