Skip to content

Commit 84644fa

Browse files
stefanozampinistefanozampini
committed
Support MATIS
1 parent 3034b59 commit 84644fa

File tree

6 files changed

+130
-69
lines changed

6 files changed

+130
-69
lines changed

‎cpp/dolfinx/fem/assembler.h‎

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -575,6 +575,29 @@ void set_diagonal(auto set_fn, std::span<const std::int32_t> rows,
575575
}
576576
}
577577

578+
/// @brief Sets values to the diagonal of a matrix for specified rows.
579+
///
580+
/// This function is typically called after assembly. The assembly
581+
/// function zeroes Dirichlet rows and columns. For block matrices, this
582+
/// function should normally be called only on the diagonal blocks, i.e.
583+
/// blocks for which the test and trial spaces are the same.
584+
///
585+
/// @param[in] set_fn The function for setting values to a matrix.
586+
/// @param[in] rows Row blocks, in local indices, for which to add a
587+
/// value to the diagonal.
588+
/// @param[in] diagonal Values to add to the diagonal for the specified
589+
/// rows.
590+
template <dolfinx::scalar T>
591+
void set_diagonal(auto set_fn, std::span<const std::int32_t> rows,
592+
std::span<const T> diagonal)
593+
{
594+
assert(diagonal.size() == rows.size());
595+
for (std::size_t i = 0; i < rows.size(); ++i)
596+
{
597+
set_fn(rows.subspan(i, 1), rows.subspan(i, 1), diagonal.subspan(i, 1));
598+
}
599+
}
600+
578601
/// @brief Sets a value to the diagonal of the matrix for rows with a
579602
/// Dirichlet boundary conditions applied.
580603
///
@@ -587,22 +610,43 @@ void set_diagonal(auto set_fn, std::span<const std::int32_t> rows,
587610
/// @param[in] set_fn The function for setting values to a matrix.
588611
/// @param[in] V The function space for the rows and columns of the
589612
/// matrix. It is used to extract only the Dirichlet boundary conditions
590-
/// that are define on V or subspaces of V.
613+
/// that are defined on V or subspaces of V.
591614
/// @param[in] bcs The Dirichlet boundary conditions.
592615
/// @param[in] diagonal Value to add to the diagonal for rows with a
593616
/// boundary condition applied.
617+
/// @param[in] unassembled Whether the matrix is in unassembled format
594618
template <dolfinx::scalar T, std::floating_point U>
595619
void set_diagonal(
596620
auto set_fn, const FunctionSpace<U>& V,
597621
const std::vector<std::reference_wrapper<const DirichletBC<T, U>>>& bcs,
598-
T diagonal = 1.0)
622+
T diagonal = 1.0, bool unassembled = false)
599623
{
600624
for (auto& bc : bcs)
601625
{
602626
if (V.contains(*bc.get().function_space()))
603627
{
604628
const auto [dofs, range] = bc.get().dof_indices();
605-
set_diagonal(set_fn, dofs.first(range), diagonal);
629+
if (unassembled)
630+
{
631+
// We need to insert on ghost indices too; since MATIS is not
632+
// designed to support the INSERT_VALUES stage, we scale ALL
633+
// the diagonal values we insert into the matrix
634+
// TODO: move scaling factor computation to DirichletBC or IndexMap?
635+
auto adjlist = bc.get().function_space()->dofmap()->index_map->index_to_dest_ranks();
636+
const auto off = adjlist.offsets();
637+
std::vector<std::int32_t> number_of_sharing_ranks(off.size());
638+
for (size_t i = 0; i < off.size() - 1; ++i)
639+
number_of_sharing_ranks[i] = off[i+1] - off[i] + 1;
640+
std::vector<T> data(dofs.size());
641+
for (size_t i = 0; i < dofs.size(); ++i)
642+
data[i] = diagonal / T(number_of_sharing_ranks[dofs[i]]);
643+
std::span<const T> data_span = std::span(data);
644+
set_diagonal(set_fn, dofs, data_span);
645+
}
646+
else
647+
{
648+
set_diagonal(set_fn, dofs.first(range), diagonal);
649+
}
606650
}
607651
}
608652
}

‎cpp/dolfinx/fem/petsc.h‎

Lines changed: 10 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -122,12 +122,6 @@ Mat create_matrix_block(
122122
la::SparsityPattern pattern(mesh->comm(), p, maps, bs_dofs);
123123
pattern.finalize();
124124

125-
// FIXME: Add option to pass customised local-to-global map to PETSc
126-
// Mat constructor
127-
128-
// Initialise matrix
129-
Mat A = la::petsc::create_matrix(mesh->comm(), pattern, type);
130-
131125
// Create row and column local-to-global maps (field0, field1, field2,
132126
// etc), i.e. ghosts of field0 appear before owned indices of field1
133127
std::array<std::vector<PetscInt>, 2> _maps;
@@ -161,29 +155,27 @@ Mat create_matrix_block(
161155
}
162156

163157
// Create PETSc local-to-global map/index sets and attach to matrix
164-
ISLocalToGlobalMapping petsc_local_to_global0;
165-
ISLocalToGlobalMappingCreate(MPI_COMM_SELF, 1, _maps[0].size(),
158+
ISLocalToGlobalMapping petsc_local_to_global0, petsc_local_to_global1;
159+
ISLocalToGlobalMappingCreate(mesh->comm(), 1, _maps[0].size(),
166160
_maps[0].data(), PETSC_COPY_VALUES,
167161
&petsc_local_to_global0);
168162
if (V[0] == V[1])
169163
{
170-
MatSetLocalToGlobalMapping(A, petsc_local_to_global0,
171-
petsc_local_to_global0);
172-
ISLocalToGlobalMappingDestroy(&petsc_local_to_global0);
164+
PetscObjectReference((PetscObject)petsc_local_to_global0);
165+
petsc_local_to_global1 = petsc_local_to_global0;
173166
}
174167
else
175168
{
176-
177-
ISLocalToGlobalMapping petsc_local_to_global1;
178-
ISLocalToGlobalMappingCreate(MPI_COMM_SELF, 1, _maps[1].size(),
169+
ISLocalToGlobalMappingCreate(mesh->comm(), 1, _maps[1].size(),
179170
_maps[1].data(), PETSC_COPY_VALUES,
180171
&petsc_local_to_global1);
181-
MatSetLocalToGlobalMapping(A, petsc_local_to_global0,
182-
petsc_local_to_global1);
183-
ISLocalToGlobalMappingDestroy(&petsc_local_to_global0);
184-
ISLocalToGlobalMappingDestroy(&petsc_local_to_global1);
185172
}
186173

174+
// Initialise matrix
175+
Mat A = la::petsc::create_matrix(mesh->comm(), pattern, type, petsc_local_to_global0, petsc_local_to_global1);
176+
ISLocalToGlobalMappingDestroy(&petsc_local_to_global0);
177+
ISLocalToGlobalMappingDestroy(&petsc_local_to_global1);
178+
187179
return A;
188180
}
189181

‎cpp/dolfinx/la/petsc.cpp‎

Lines changed: 52 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -232,7 +232,9 @@ void la::petsc::scatter_local_vectors(
232232
}
233233
//-----------------------------------------------------------------------------
234234
Mat la::petsc::create_matrix(MPI_Comm comm, const SparsityPattern& sp,
235-
std::optional<std::string> type)
235+
std::optional<std::string> type,
236+
std::optional<ISLocalToGlobalMapping> rlgmap,
237+
std::optional<ISLocalToGlobalMapping> clgmap)
236238
{
237239
PetscErrorCode ierr;
238240
Mat A;
@@ -289,63 +291,75 @@ Mat la::petsc::create_matrix(MPI_Comm comm, const SparsityPattern& sp,
289291
_nnz_offdiag[i] = bs[1] * sp.nnz_off_diag(i / bs[0]);
290292
}
291293

292-
// Allocate space for matrix
293-
ierr = MatXAIJSetPreallocation(A, _bs, _nnz_diag.data(), _nnz_offdiag.data(),
294-
nullptr, nullptr);
295-
if (ierr != 0)
296-
petsc::error(ierr, __FILE__, "MatXIJSetPreallocation");
297-
298-
// Set block sizes
299-
ierr = MatSetBlockSizes(A, bs[0], bs[1]);
300-
if (ierr != 0)
301-
petsc::error(ierr, __FILE__, "MatSetBlockSizes");
302-
303294
// Create PETSc local-to-global map/index sets
304-
ISLocalToGlobalMapping local_to_global0;
305-
const std::vector map0 = maps[0]->global_indices();
306-
const std::vector<PetscInt> _map0(map0.begin(), map0.end());
307-
ierr = ISLocalToGlobalMappingCreate(MPI_COMM_SELF, bs[0], _map0.size(),
308-
_map0.data(), PETSC_COPY_VALUES,
309-
&local_to_global0);
310-
311-
if (ierr != 0)
312-
petsc::error(ierr, __FILE__, "ISLocalToGlobalMappingCreate");
295+
ISLocalToGlobalMapping local_to_global0, local_to_global1 = NULL;
296+
if (rlgmap)
297+
{
298+
ierr = PetscObjectReference((PetscObject)rlgmap.value());
299+
if (ierr != 0)
300+
petsc::error(ierr, __FILE__, "PetscObjectReference");
301+
local_to_global0 = rlgmap.value();
302+
}
303+
else
304+
{
305+
const std::vector map0 = maps[0]->global_indices();
306+
const std::vector<PetscInt> _map0(map0.begin(), map0.end());
307+
ierr = ISLocalToGlobalMappingCreate(comm, bs[0], _map0.size(),
308+
_map0.data(), PETSC_COPY_VALUES,
309+
&local_to_global0);
310+
if (ierr != 0)
311+
petsc::error(ierr, __FILE__, "ISLocalToGlobalMappingCreate");
312+
}
313313

314314
// Check for common index maps
315-
if (maps[0] == maps[1] and bs[0] == bs[1])
315+
if (maps[0] == maps[1] and bs[0] == bs[1] and !clgmap)
316316
{
317317
ierr = MatSetLocalToGlobalMapping(A, local_to_global0, local_to_global0);
318318
if (ierr != 0)
319319
petsc::error(ierr, __FILE__, "MatSetLocalToGlobalMapping");
320320
}
321321
else
322322
{
323-
ISLocalToGlobalMapping local_to_global1;
324-
const std::vector map1 = maps[1]->global_indices();
325-
const std::vector<PetscInt> _map1(map1.begin(), map1.end());
326-
ierr = ISLocalToGlobalMappingCreate(MPI_COMM_SELF, bs[1], _map1.size(),
327-
_map1.data(), PETSC_COPY_VALUES,
328-
&local_to_global1);
329-
if (ierr != 0)
330-
petsc::error(ierr, __FILE__, "ISLocalToGlobalMappingCreate");
323+
if (clgmap)
324+
{
325+
ierr = PetscObjectReference((PetscObject)clgmap.value());
326+
if (ierr != 0)
327+
petsc::error(ierr, __FILE__, "PetscObjectReference");
328+
local_to_global1 = clgmap.value();
329+
}
330+
else
331+
{
332+
const std::vector map1 = maps[1]->global_indices();
333+
const std::vector<PetscInt> _map1(map1.begin(), map1.end());
334+
ierr = ISLocalToGlobalMappingCreate(comm, bs[1], _map1.size(),
335+
_map1.data(), PETSC_COPY_VALUES,
336+
&local_to_global1);
337+
if (ierr != 0)
338+
petsc::error(ierr, __FILE__, "ISLocalToGlobalMappingCreate");
339+
}
331340
ierr = MatSetLocalToGlobalMapping(A, local_to_global0, local_to_global1);
332341
if (ierr != 0)
333342
petsc::error(ierr, __FILE__, "MatSetLocalToGlobalMapping");
334-
ierr = ISLocalToGlobalMappingDestroy(&local_to_global1);
335-
if (ierr != 0)
336-
petsc::error(ierr, __FILE__, "ISLocalToGlobalMappingDestroy");
337343
}
338344

339345
// Clean up local-to-global 0
340346
ierr = ISLocalToGlobalMappingDestroy(&local_to_global0);
347+
if (ierr != 0)
348+
petsc::error(ierr, __FILE__, "ISLocalToGlobalMappingDestroy");
349+
ierr = ISLocalToGlobalMappingDestroy(&local_to_global1);
341350
if (ierr != 0)
342351
petsc::error(ierr, __FILE__, "ISLocalToGlobalMappingDestroy");
343352

344-
// Note: This should be called after having set the local-to-global
345-
// map for MATIS (this is a dummy call if A is not of type MATIS)
346-
// ierr = MatISSetPreallocation(A, 0, _nnz_diag.data(), 0,
347-
// _nnz_offdiag.data()); if (ierr != 0)
348-
// error(ierr, __FILE__, "MatISSetPreallocation");
353+
// Allocate space for matrix
354+
ierr = MatXAIJSetPreallocation(A, _bs, _nnz_diag.data(), _nnz_offdiag.data(),
355+
nullptr, nullptr);
356+
if (ierr != 0)
357+
petsc::error(ierr, __FILE__, "MatXIJSetPreallocation");
358+
359+
// Set block sizes
360+
ierr = MatSetBlockSizes(A, bs[0], bs[1]);
361+
if (ierr != 0)
362+
petsc::error(ierr, __FILE__, "MatSetBlockSizes");
349363

350364
// Set some options on Mat object
351365
ierr = MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);

‎cpp/dolfinx/la/petsc.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,9 @@ void scatter_local_vectors(
120120
/// Create a PETSc Mat. Caller is responsible for destroying the
121121
/// returned object.
122122
Mat create_matrix(MPI_Comm comm, const SparsityPattern& sp,
123-
std::optional<std::string> type = std::nullopt);
123+
std::optional<std::string> type = std::nullopt,
124+
std::optional<ISLocalToGlobalMapping> rlgmap = std::nullopt,
125+
std::optional<ISLocalToGlobalMapping> clgmap = std::nullopt);
124126

125127
/// Create PETSc MatNullSpace. Caller is responsible for destruction
126128
/// returned object.

‎python/dolfinx/fem/petsc.py‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -513,13 +513,14 @@ def _(
513513
A.assemble(PETSc.Mat.AssemblyType.FLUSH) # type: ignore[attr-defined]
514514

515515
# Set diagonal
516-
for i, a_row in enumerate(a):
517-
for j, a_sub in enumerate(a_row):
518-
if a_sub is not None:
519-
Asub = A.getLocalSubMatrix(is0[i], is1[j])
520-
if a_sub.function_spaces[0] is a_sub.function_spaces[1]:
521-
_cpp.fem.petsc.insert_diagonal(Asub, a_sub.function_spaces[0], _bcs, diag)
522-
A.restoreLocalSubMatrix(is0[i], is1[j], Asub)
516+
if len(_bcs):
517+
for i, a_row in enumerate(a):
518+
for j, a_sub in enumerate(a_row):
519+
if a_sub is not None:
520+
if a_sub.function_spaces[0] is a_sub.function_spaces[1]:
521+
Asub = A.getLocalSubMatrix(is0[i], is1[j])
522+
_cpp.fem.petsc.insert_diagonal(Asub, a_sub.function_spaces[0], _bcs, diag)
523+
A.restoreLocalSubMatrix(is0[i], is1[j], Asub)
523524
else: # Non-blocked
524525
constants = pack_constants(a) if constants is None else constants # type: ignore[assignment]
525526
coeffs = pack_coefficients(a) if coeffs is None else coeffs # type: ignore[assignment]
@@ -528,7 +529,8 @@ def _(
528529
if a.function_spaces[0] is a.function_spaces[1]:
529530
A.assemblyBegin(PETSc.Mat.AssemblyType.FLUSH) # type: ignore[attr-defined]
530531
A.assemblyEnd(PETSc.Mat.AssemblyType.FLUSH) # type: ignore[attr-defined]
531-
_cpp.fem.petsc.insert_diagonal(A, a.function_spaces[0], _bcs, diag)
532+
if len(_bcs):
533+
_cpp.fem.petsc.insert_diagonal(A, a.function_spaces[0], _bcs, diag)
532534

533535
return A
534536

@@ -849,6 +851,8 @@ def __init__(
849851

850852
# For nest matrices kind can be a nested list.
851853
kind = "nest" if self.A.getType() == PETSc.Mat.Type.NEST else kind # type: ignore[attr-defined]
854+
if kind == "is":
855+
kind = "mpi"
852856
assert kind is None or isinstance(kind, str)
853857
self._b = create_vector(self.L, kind=kind)
854858
self._x = create_vector(self.L, kind=kind)
@@ -1326,6 +1330,8 @@ def __init__(
13261330

13271331
# Determine the vector kind based on the matrix type
13281332
kind = "nest" if self._A.getType() == PETSc.Mat.Type.NEST else kind # type: ignore[attr-defined]
1333+
if kind == "is":
1334+
kind = "mpi"
13291335
assert kind is None or isinstance(kind, str)
13301336
self._b = create_vector(self.F, kind=kind)
13311337
self._x = create_vector(self.F, kind=kind)

‎python/dolfinx/wrappers/petsc.cpp‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -393,9 +393,12 @@ void petsc_fem_module(nb::module_& m)
393393
_bcs.push_back(*bc);
394394
}
395395

396+
PetscBool flg;
397+
PetscObjectTypeCompare((PetscObject)A, MATIS, &flg);
398+
396399
dolfinx::fem::set_diagonal(
397400
dolfinx::la::petsc::Matrix::set_fn(A, INSERT_VALUES), V, _bcs,
398-
diagonal);
401+
diagonal, flg ? true : false);
399402
},
400403
nb::arg("A"), nb::arg("V"), nb::arg("bcs"), nb::arg("diagonal"));
401404

0 commit comments

Comments
 (0)