DOLFINx 0.12.0.0
DOLFINx C++
Loading...
Searching...
No Matches
matmul.h
1// Copyright (C) 2026 Chris Richardson
2// Copyright (C) 2026 Garth N. Wells
3//
4// This file is part of DOLFINx (https://www.fenicsproject.org)
5//
6// SPDX-License-Identifier: LGPL-3.0-or-later
7
8#pragma once
9
10#include <dolfinx/common/IndexMap.h>
11#include <dolfinx/common/log.h>
12#include <stdexcept>
13
14#include <dolfinx/la/MatrixCSR.h>
15#include <mpi.h>
16
17#include <algorithm>
18#include <numeric>
19#include <span>
20#include <vector>
21
22namespace dolfinx::la
23{
24
38namespace impl
39{
40
53{
55 std::shared_ptr<const common::IndexMap> _row_map;
56
60 std::shared_ptr<const common::IndexMap> _col_map;
61
64 std::span<const std::int32_t> _cols;
65
68 std::span<const std::int64_t> _offsets;
69
76 std::span<const std::int32_t> _off_diag;
77
79 std::array<int, 2> _bs;
80
82 std::shared_ptr<const common::IndexMap> index_map(int dim) const
83 {
84 return dim == 0 ? _row_map : _col_map;
85 }
86
88 int block_size(int i) const { return _bs[i]; }
89
91 std::pair<std::span<const std::int32_t>, std::span<const std::int64_t>>
92 graph() const
93 {
94 return {_cols, _offsets};
95 }
96
98 std::span<const std::int32_t> off_diagonal_offsets() const
99 {
100 return _off_diag;
101 }
102};
103
115template <typename T, int BS = -1>
116std::tuple<std::shared_ptr<common::IndexMap>, std::vector<std::int32_t>,
117 std::vector<std::int32_t>, std::vector<T>>
120{
121 // The ghost columns of A are global row indices into B.
122 auto col_map_A = A.index_map(1); // column IndexMap of A
123 auto row_map_B = B.index_map(0); // row IndexMap of B
124
125 // Runtime block size: use compile-time BS when provided, else read from B.
126 const int bs_rt = (BS < 0) ? B.block_size()[0] : BS;
127 const int bsq = bs_rt * bs_rt;
128
129 // Create neighborhood comms for col_map_A
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();
133
134 // Serial fast-path: with a single rank there are no ghost columns and
135 // no MPI communication to perform.
136 // In parallel, even ranks without neighbours must participate in the
137 // communicator creation and IndexMap collectives below.
138 // Use col_map_B — not col_map_A — for new_col_map so that any existing
139 // ghost columns in B are preserved in the product's column map.
140
141 int comm_size = dolfinx::MPI::size(comm);
142 if (comm_size == 1)
143 {
144 auto col_map_B = B.index_map(1);
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>{}};
150 }
151
152 // src is guaranteed sorted (Scatterer asserts this), so use binary search
153 // rather than a heap-allocated map.
154 auto rank_to_nbr = [&src](int r) -> int
155 {
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));
159 };
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);
168
169 // Global indices of the rows of B we need (= ghost cols of A)
170 std::span<const std::int64_t> required_globals_map = col_map_A->ghosts();
171
172 // Owning rank for each ghost col of A (= source rank for each needed row)
173 std::span<const int> ghost_owners = col_map_A->owners();
174
175 std::vector<int> send_count(src.size(), 0);
176 std::vector<int> recv_count(dest.size(), 0);
177 send_count.reserve(1); // ensure data is not a nullptr
178 recv_count.reserve(1); // ensure data is not a nullptr
179
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);
184
185 // Send and recv displacements
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()));
192
193 // perm[i] = position in the send buffer (and received data) of original
194 // ghost column i, so we can invert the ordering after receiving.
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)
198 {
199 int pos = rank_to_nbr(ghost_owners[i]);
200 perm[i] = send_disp[pos]; // position before increment
201 required_globals[send_disp[pos]++] = required_globals_map[i];
202 }
203
204 // Reset send_disp
205 send_disp[0] = 0;
206 std::partial_sum(send_count.begin(), send_count.end(),
207 std::next(send_disp.begin()));
208
209 // Send the global row indices we need; receive the indices others need from
210 // us send buffer = required_globals (already ordered by src[])
211 std::vector<std::int64_t> recv_row_globals(recv_disp.back());
212
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,
216 neigh_comm_fwd);
217
218 // Convert received global row indices to local row indices in B.
219 // The received indices are global rows of B that remote ranks need.
220 // They must all be owned by this rank (by construction).
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);
223
224 // Compute size of send_data and communicate
225 auto row_ptr_B = B.row_ptr();
226
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); // ensure data is not a nullptr
232 recv_entry_count.reserve(1); // ensure data is not a nullptr
233 assert(recv_count.size() == dest.size());
234 for (std::size_t p = 0; p < recv_count.size(); ++p)
235 {
236 for (std::int32_t i = recv_disp[p]; i < recv_disp[p + 1]; ++i)
237 {
238 int r = recv_row_locals[i];
239 send_entry_count[p] += row_ptr_B[r + 1] - row_ptr_B[r];
240 }
241 }
242 std::partial_sum(send_entry_count.begin(), send_entry_count.end(),
243 std::next(send_entry_disp.begin()));
244
245 MPI_Neighbor_alltoall(send_entry_count.data(), 1, MPI_INT,
246 recv_entry_count.data(), 1, MPI_INT, neigh_comm_rev);
247
248 std::partial_sum(recv_entry_count.begin(), recv_entry_count.end(),
249 std::next(recv_entry_disp.begin()));
250
251 // Pack and send data values.
252 // send_entry_count/disp count graph entries (CSR nodes); multiply by bsq
253 // for the value count since each entry holds a BS×BS block.
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);
257 std::size_t k = 0;
258 for (int r : recv_row_locals)
259 {
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));
264 k += row_size * bsq;
265 }
266
267 // Communicate values using a compound MPI datatype of bsq scalars per
268 // graph entry, so count/disp arrays index graph entries directly.
269 const dolfinx::MPI::Datatype<T> block_type(bsq);
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);
275
276 // Pack and send column indices
277 auto cols_B = B.cols();
278 auto col_map_B = B.index_map(1);
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();
281
282 // Send column indices (global indexing) and their owner ranks
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());
287
288 // Send row sizes so receivers can reconstruct the CSR row_ptr
289 std::vector<int> send_row_size(recv_disp.back());
290 std::vector<int> recv_row_size(send_disp.back());
291
292 k = 0;
293 int rank = dolfinx::MPI::rank(col_map_A->comm());
294 for (std::size_t p = 0; p < dest.size(); ++p)
295 {
296 for (std::int32_t i = recv_disp[p]; i < recv_disp[p + 1]; ++i)
297 {
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)
304 {
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)
308 ? rank
309 : ghost_owners_B[local_col - num_owned_cols_B];
310 }
311 k += row_size;
312 send_row_size[i] = row_size;
313 }
314 }
315
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,
327 neigh_comm_rev);
328
329 MPI_Comm_free(&neigh_comm_fwd);
330 MPI_Comm_free(&neigh_comm_rev);
331
332 // Build CSR row pointer for the received ghost rows.
333 // recv_vals and recv_col_data are the values and (global) column indices;
334 // recv_row_size[i] is the number of entries in ghost row i.
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()));
338
339 // ghost_row_ptr, recv_vals, recv_col_data now form a CSR structure for
340 // the ghost rows of B needed by this rank.
341
342 // Collect ghost column indices with their owners.
343 // Start from the received ghost rows, then add existing ghosts of col_map_B.
344 // Exclude anything in the locally owned range of col_map_B, then deduplicate.
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; };
348
349 // Merge received (global_col, owner) pairs with the existing ghosts.
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]});
355
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]});
360
361 // Sort by global index and deduplicate.
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());
366
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)
370 {
371 unique_cols[i] = col_owner_pairs[i].first;
372 unique_col_owners[i] = col_owner_pairs[i].second;
373 }
374
375 // Build the extended column IndexMap for B.
376 auto new_col_map = std::make_shared<dolfinx::common::IndexMap>(
377 comm, col_map_B->size_local(), unique_cols, unique_col_owners);
378
379 // Convert received global column indices to local indices using new_col_map.
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);
382
383 // The received data is in send-buffer order (grouped by src rank).
384 // Reorder it back to the original ghost column order of col_map_A so that
385 // sparsity_pattern and matmul can index by (j - num_local_rows_B).
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());
392
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)
396 {
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));
406 }
407
408 return {new_col_map, orig_ghost_row_ptr, orig_recv_col_local, orig_recv_vals};
409}
410
432template <typename T, int BS = -1>
433std::tuple<std::vector<std::int64_t>, std::vector<std::int32_t>,
434 std::vector<std::int32_t>, std::vector<T>>
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)
439{
440 // Runtime block size: use compile-time BS when provided, else read from A.
441 const int bs_rt = (BS < 0) ? A.block_size()[0] : BS;
442 const int bsq = bs_rt * bs_rt;
443
444 auto col_map_B = B.index_map(1);
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();
449
450 auto row_ptr_A = A.row_ptr();
451 auto offdiag_ptr_A = A.off_diag_offset();
452 auto cols_A = A.cols();
453 auto vals_A = A.values();
454 auto row_ptr_B = B.row_ptr();
455 auto cols_B = B.cols();
456 auto vals_B = B.values();
457
458 // Remap B's ghost column indices into new_col_map's index space.
459 // Owned column indices (0..num_owned_cols_B-1) are identical in both maps.
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);
463
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());
467
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;
474
475 // Dense accumulator: acc[k*bsq..(k+1)*bsq-1] holds the running BS×BS block
476 // sum for output column (node) k. in_row[k] marks whether column k was
477 // touched in the current row. Both are reset via row_cols at the end of each
478 // row (O(nnz/row), not O(num_cols_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;
482
483 // Block GEMM: C[r,s] += sum_t A[r,t] * B[t,s] (row-major BS×BS).
484 // When BS >= 1 the loop bounds are compile-time constants, enabling
485 // unrolling; when BS < 0 the runtime size bs_rt is used instead.
486 auto gemm_block = [&](const T* a_blk, const T* b_blk, T* c_blk)
487 {
488 if constexpr (BS < 0)
489 {
490 for (int r = 0; r < bs_rt; ++r)
491 for (int t = 0; t < bs_rt; ++t)
492 {
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];
496 }
497 }
498 else
499 {
500 for (int r = 0; r < BS; ++r)
501 for (int t = 0; t < BS; ++t)
502 {
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];
506 }
507 }
508 };
509
510 for (std::int32_t i = 0; i < num_local_rows_A; ++i)
511 {
512 row_cols.clear();
513
514 // Local columns of A → local rows of B.
515 for (std::int32_t ka = row_ptr_A[i]; ka < offdiag_ptr_A[i]; ++ka)
516 {
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)
520 {
521 const std::int32_t c = cols_B[kb];
522 const std::int32_t k
523 = c < num_owned_cols_B ? c : b_ghost_remap[c - num_owned_cols_B];
524 if (!in_row[k])
525 {
526 in_row[k] = true;
527 row_cols.push_back(k);
528 }
529 gemm_block(a_blk, vals_B.data() + kb * bsq, acc.data() + k * bsq);
530 }
531 }
532
533 // Ghost columns of A → ghost rows of B.
534 for (std::int32_t ka = offdiag_ptr_A[i]; ka < row_ptr_A[i + 1]; ++ka)
535 {
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)
540 {
541 const std::int32_t k = ghost_cols[kb];
542 if (!in_row[k])
543 {
544 in_row[k] = true;
545 row_cols.push_back(k);
546 }
547 gemm_block(a_blk, ghost_vals.data() + kb * bsq, acc.data() + k * bsq);
548 }
549 }
550
551 // Remove nodes whose accumulated block is identically zero (exact
552 // cancellation in the SpGEMM).
553 auto pos = row_cols.begin();
554 while (pos != row_cols.end())
555 {
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));
560 if (all_zero)
561 {
562 in_row[*pos] = false;
563 pos = row_cols.erase(pos);
564 }
565 else
566 ++pos;
567 }
568
569 // Sort touched indices to satisfy the CSR sorted-column invariant
570 // required by off_diag_offset and MatrixCSR.
571 std::sort(row_cols.begin(), row_cols.end());
572
573 // Diagonal block boundary: count entries in the owned column range.
574 // row_cols is sorted, so a single lower_bound suffices.
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()));
578
579 // Flush accumulator to output and reset both acc and in_row.
580 for (const std::int32_t c : row_cols)
581 {
582 cols_C.push_back(c);
583 T* c_blk = acc.data() + c * bsq;
584 for (int q = 0; q < bsq; ++q)
585 {
586 vals_C.push_back(c_blk[q]);
587 c_blk[q] = T(0);
588 }
589 in_row[c] = false;
590 }
591 row_ptr_C.push_back(static_cast<std::int64_t>(cols_C.size()));
592 }
593
594 return {std::move(row_ptr_C), std::move(off_diag_offsets_C),
595 std::move(cols_C), std::move(vals_C)};
596}
597
598} // namespace impl
599
615template <typename T, int BS = -1>
618{
619 dolfinx::common::Timer t_spgemm("MatrixCSR SpGEMM");
620
621 // Validate block sizes.
622 if (A.block_size() != B.block_size())
623 {
624 throw std::invalid_argument(
625 "matmul: A and B must have the same block size");
626 }
627 if (A.block_size()[0] != A.block_size()[1])
628 {
629 throw std::invalid_argument(
630 "matmul: only square block sizes are supported");
631 }
632 if constexpr (BS >= 1)
633 {
634 if (A.block_size() != std::array<int, 2>{BS, BS})
635 {
636 throw std::invalid_argument(
637 "matmul: template parameter BS does not match matrix block size");
638 }
639 }
640
641 const std::array<int, 2> bs = A.block_size();
642
643 // Fetch ghost rows of B needed to multiply against A's ghost columns.
644 logger().debug("Fetch remote rows of B in C=A*B");
645 auto [new_col_map, ghost_row_ptr, ghost_cols, ghost_vals]
646 = impl::fetch_ghost_rows<T, BS>(A, B);
647
648 // Single pass: compute sparsity, off-diagonal boundaries, and values.
649 logger().debug("Compute sparsity and values of C=A*B");
650 auto [C_row_ptr, C_off_diag_offsets, C_cols, C_vals_vec]
651 = impl::matmul<T, BS>(A, B, new_col_map,
652 std::span<const std::int32_t>(ghost_row_ptr),
653 std::span<const std::int32_t>(ghost_cols),
654 std::span<const T>(ghost_vals));
655
656 auto C_row_map = std::make_shared<common::IndexMap>(
657 A.index_map(0)->comm(), A.index_map(0)->size_local());
658 impl::Sparsity sp{C_row_map, new_col_map, C_cols,
659 C_row_ptr, C_off_diag_offsets, bs};
661 std::copy(C_vals_vec.begin(), C_vals_vec.end(), C.values().begin());
662
663 return C;
664}
665
666} // namespace dolfinx::la
An MPI datatype for count contiguous values of type T, and manage its lifetime.
Definition MPI.h:346
MPI_Datatype type() const noexcept
The datatype to pass to MPI.
Definition MPI.h:395
Timer for measuring and logging elapsed time durations.
Definition Timer.h:41
Distributed sparse matrix using compressed sparse row storage.
Definition MatrixCSR.h:74
std::shared_ptr< const common::IndexMap > index_map(int dim) const
Index map for the row or column space.
Definition MatrixCSR.h:574
const rowptr_container_type & off_diag_offset() const
Get the start of off-diagonal (unowned columns) on each row, allowing the matrix to be split (virtual...
Definition MatrixCSR.h:604
container_type & values()
Get local data values.
Definition MatrixCSR.h:581
const column_container_type & cols() const
Definition MatrixCSR.h:593
std::array< int, 2 > block_size() const
Get block sizes.
Definition MatrixCSR.h:611
const rowptr_container_type & row_ptr() const
Get local row pointers.
Definition MatrixCSR.h:589
int size(MPI_Comm comm)
Definition MPI.cpp:81
int rank(MPI_Comm comm)
Return process rank for the communicator.
Definition MPI.cpp:73
Fetch the rows of B that correspond to the ghost columns of A.
Definition matmul.h:39
std::tuple< std::shared_ptr< common::IndexMap >, std::vector< std::int32_t >, std::vector< std::int32_t >, std::vector< T > > fetch_ghost_rows(const dolfinx::la::MatrixCSR< T > &A, const dolfinx::la::MatrixCSR< T > &B)
Fetch the rows of Matrix B which are referenced by the ghost columns of Matrix A.
Definition matmul.h:118
std::tuple< std::vector< std::int64_t >, std::vector< std::int32_t >, std::vector< std::int32_t >, std::vector< T > > matmul(const dolfinx::la::MatrixCSR< T > &A, const dolfinx::la::MatrixCSR< T > &B, std::shared_ptr< const common::IndexMap > new_col_map, std::span< const std::int32_t > ghost_row_ptr, std::span< const std::int32_t > ghost_cols, std::span< const T > ghost_vals)
Compute the sparsity pattern and values of C = A*B in a single pass.
Definition matmul.h:435
Linear algebra interface.
Definition dolfinx_la.h:7
dolfinx::la::MatrixCSR< T > matmul(const dolfinx::la::MatrixCSR< T > &A, const dolfinx::la::MatrixCSR< T > &B)
Compute C = A*B as a distributed MatrixCSR.
Definition matmul.h:616
auto logger() -> spdlog::logger &
Dolfinx's logger.
Definition log.cpp:24
Lightweight sparsity descriptor satisfying the SparsityImplementation concept required by the MatrixC...
Definition matmul.h:53
std::shared_ptr< const common::IndexMap > index_map(int dim) const
Return the row (dim == 0) or column (dim == 1) IndexMap.
Definition matmul.h:82
int block_size(int i) const
Return the block size.
Definition matmul.h:88
std::span< const std::int32_t > off_diagonal_offsets() const
Return the per-row diagonal block sizes (see _off_diag).
Definition matmul.h:98
std::pair< std::span< const std::int32_t >, std::span< const std::int64_t > > graph() const
Return the CSR graph as (column_indices, row_pointers).
Definition matmul.h:92