Initial support for PETSc MATIS matrix format #3885
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -575,6 +575,29 @@ void set_diagonal(auto set_fn, std::span<const std::int32_t> rows, | |
| } | ||
| } | ||
|
|
||
| /// @brief Sets values to the diagonal of a matrix for specified rows. | ||
| /// | ||
| /// This function is typically called after assembly. The assembly | ||
| /// function zeroes Dirichlet rows and columns. For block matrices, this | ||
| /// function should normally be called only on the diagonal blocks, i.e. | ||
| /// blocks for which the test and trial spaces are the same. | ||
| /// | ||
| /// @param[in] set_fn The function for setting values to a matrix. | ||
| /// @param[in] rows Row blocks, in local indices, for which to add a | ||
| /// value to the diagonal. | ||
| /// @param[in] diagonal Values to add to the diagonal for the specified | ||
| /// rows. | ||
| template <dolfinx::scalar T> | ||
| void set_diagonal(auto set_fn, std::span<const std::int32_t> rows, | ||
| std::span<const T> diagonal) | ||
| { | ||
| assert(diagonal.size() == rows.size()); | ||
| for (std::size_t i = 0; i < rows.size(); ++i) | ||
| { | ||
| set_fn(rows.subspan(i, 1), rows.subspan(i, 1), diagonal.subspan(i, 1)); | ||
| } | ||
| } | ||
|
|
||
| /// @brief Sets a value to the diagonal of the matrix for rows with a | ||
| /// Dirichlet boundary conditions applied. | ||
| /// | ||
|
|
@@ -587,22 +610,43 @@ void set_diagonal(auto set_fn, std::span<const std::int32_t> rows, | |
| /// @param[in] set_fn The function for setting values to a matrix. | ||
| /// @param[in] V The function space for the rows and columns of the | ||
| /// matrix. It is used to extract only the Dirichlet boundary conditions | ||
| /// that are defined on V or subspaces of V. | ||
| /// @param[in] bcs The Dirichlet boundary conditions. | ||
| /// @param[in] diagonal Value to add to the diagonal for rows with a | ||
| /// boundary condition applied. | ||
| /// @param[in] unassembled Whether the matrix is in unassembled format | ||
| template <dolfinx::scalar T, std::floating_point U> | ||
| void set_diagonal( | ||
| auto set_fn, const FunctionSpace<U>& V, | ||
| const std::vector<std::reference_wrapper<const DirichletBC<T, U>>>& bcs, | ||
| T diagonal = 1.0, bool unassembled = false) | ||
| { | ||
| for (auto& bc : bcs) | ||
| { | ||
| if (V.contains(*bc.get().function_space())) | ||
| { | ||
| const auto [dofs, range] = bc.get().dof_indices(); | ||
| if (unassembled) | ||
| { | ||
| // We need to insert on ghost indices too; since MATIS is not | ||
| // designed to support the INSERT_VALUES stage, we scale ALL | ||
| // the diagonal values we insert into the matrix | ||
| // TODO: move scaling factor computation to DirichletBC or IndexMap? | ||
| auto adjlist = bc.get().function_space()->dofmap()->index_map->index_to_dest_ranks(); | ||
stefanozampini marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| const auto off = adjlist.offsets(); | ||
| std::vector<std::int32_t> number_of_sharing_ranks(off.size()); | ||
| for (size_t i = 0; i < off.size() - 1; ++i) | ||
| number_of_sharing_ranks[i] = off[i+1] - off[i] + 1; | ||
|
Comment on lines
+635
to
+639
Author
There was a problem hiding this comment. Ideally, this setup code should be moved to
Member
There was a problem hiding this comment. Hi Stefano,
that would give you an array of the number of shared procs per dof
Member
There was a problem hiding this comment. This could probably be a vector we compute once inside the DirichletBC class. What do you think @garth-wells ?
Author
There was a problem hiding this comment. @jorgensd Can you paste some code here to do it? I will wait for @garth-wells confirmation and move the setup of the counting vector in the
Member
There was a problem hiding this comment. I’m without a computer until Sunday. @IgorBaratta could you do me a solid and make a simple example for Stefano?:)
Member
There was a problem hiding this comment. We need a way to divide the insertion of 1 on the diagonal by the number of processes that has the dof, as MATIS cannot use insert, and has to add the values. I.e. if a DirichletBC dof is shared between N processes, we assign 1/N on each process. I suggest we use a vector, stored in the DirichletBC to keep track of this, a illustrated with the snippet above.
Author
There was a problem hiding this comment. @jorgensd summarized it properly. If you want to keep the same workflow for assembled and unassembled matrices, the counting vector is the only way to realize it. I think such information should belong to
Member
There was a problem hiding this comment. @stefanozampini @garth-wells I've made an attempt at this at: https://github.com/FEniCS/dolfinx/tree/dokken/suggested-matis
Member
There was a problem hiding this comment. A likely good approach is consistent handling of the LHS and RHS, i.e. use accumulate rather than set for the bc on the RHS. We did this is legacy DOLFIN. I'd need to look closely at the code for how this might work with the latest design.
Member
There was a problem hiding this comment. I think it should be quite straightforward with the approach suggested in my branch above, as it lets the user choose the insertion mode of the diagonal (which can trivially be extended to set_bc), as the DirichletBC itself holds the information about the count. |
||
| std::vector<T> data(dofs.size()); | ||
| for (size_t i = 0; i < dofs.size(); ++i) | ||
| data[i] = diagonal / T(number_of_sharing_ranks[dofs[i]]); | ||
| std::span<const T> data_span = std::span(data); | ||
| set_diagonal(set_fn, dofs, data_span); | ||
| } | ||
| else | ||
| { | ||
| set_diagonal(set_fn, dofs.first(range), diagonal); | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks problematic - quite a few types can be cast to
Tandbool, very easy for a user to make an error.Probably better to make
unassembledan enum. Easier to read too.Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Having an enum for a single function seems overkill to me.
If there's a way to extract back the
Matfromset_fn, we could test if it is aMATISinside the function, and we won't need the extra optional argument. @garth-wells what do you think?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In the branch I made, I use an insertion mode with
addorinsertto distinguish between the two, which resonates well with thescatter_fwd/scatter_reversemodes we have in the Python-interface for vectors.Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I looked at your branch. I don't think it is correct for
MATAIJ.set_diagonaluses (correctly)set_fn = dolfinx::la::petsc::Matrix::set_fn(A, INSERT_VALUES), and the values won't be accumulated intoMATAIJ. I think it is best to use an optional boolean that clearly indicates we want to insert into an unassembled (MATIS) matrix. All these problems appear because the interface forset_diagonalonly takesset_fnand not the matrix directlyThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Alternatively, use add_values formalism everywhere in matrix assembly (above all this is FEM :-)). The optimization you get by not communicating the values to be inserted into the diagonal is minimal @garth-wells ?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Right, but if we modify the input to the
set_fnto then useADD_VALUESit would resolve the issue, right?i.e.
PetscBool flg; PetscObjectTypeCompare((PetscObject)A, MATIS, &flg); dolfinx::la::VectorInsertMode mode = flg ? dolfinx::la::VectorInsertMode::add : dolfinx::la::VectorInsertMode::insert; InsertMode petsc_mode = flg ? ADD_VALUES : INSERT_VALUES; dolfinx::fem::set_diagonal( dolfinx::la::petsc::Matrix::set_fn(A, petsc_mode), V, _bcs, diagonal); diagonal, mode); },Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thinking more about it, I kind of agree that we should only have one approach, which preferably is the
add-mode, as the approach of assuming that two independent input arguments are set correctly by the user is likely to lead to a lot of confusion.Having only a single, accumulate data from all processes as
1/num_shared_procsseems sensible.The reason for DOLFINx not interfacing directly with the PETSc matrix, is that it allows us to have a single code path for multiple matrix types, like our own, built in matrices, and potentially TRILLINOS, without having to change the
set_diagonal,dolfinx::fem::assemble_*functions.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm fine with any solution you think is reasonable. Note that if you switch back to add values, you won't need to flush assembly anymore in the fem callbacks. And you can still use the optimization of adding only to diagonal entries in AIJ format (since the value in the diagonal before the call is 0.0, so
addorinsertis the same)If you want equivalent behaviour, it should be the opposite way, i.e.
petsc_mode = !flg ? ADD_VALUES : INSERT_VALUES. This is because MATIS inserts directly into the local matrix. matrix-vector products in MATIS are done via R^t A R, where R restricts from global to local, A does the local multiplications, and R^T sums up the local contributionsThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Right, I got that the wrong way around.