la_policy.hh

Go to the documentation of this file.

#pragma once
 
// external libraries
#include <deal.II/base/index_set.h>
#include <deal.II/lac/dynamic_sparsity_pattern.h>
#include <deal.II/lac/full_matrix.h>
#include <deal.II/lac/sparsity_tools.h>
 
// DiFfRG
#include <DiFfRG/common/linear_algebra.hh>
#include <DiFfRG/common/mpi.hh>
#include <DiFfRG/discretization/common/solution_view.hh>
 
// std
#include <vector>
 
namespace DiFfRG
{
  template <typename VectorType>
  void reinit_la_vector(VectorType &vec, const dealii::IndexSet &locally_owned, MPI_Comm comm)
  {
    if constexpr (is_distributed_la<VectorType>) {
      vec.reinit(locally_owned, comm);
    } else {
      (void)comm;
      vec.reinit(locally_owned.size());
    }
  }
 
  template <typename SparseMatrixType>
  void finalize_la_sparsity(dealii::DynamicSparsityPattern &dsp, get_type::SparsityPattern<SparseMatrixType> &pattern,
                            const dealii::IndexSet &locally_owned, const dealii::IndexSet &locally_relevant,
                            MPI_Comm comm)
  {
    if constexpr (is_distributed_la<SparseMatrixType>) {
      // Guarded, not merely discarded: SparsityTools::distribute_sparsity_pattern only exists in
      // an MPI-enabled deal.II, and the call's arguments are non-dependent, so an `if constexpr`
      // alone would still fail name lookup in a serial build.
#ifdef DEAL_II_WITH_MPI
      dealii::SparsityTools::distribute_sparsity_pattern(dsp, locally_owned, comm, locally_relevant);
 
      // Emphatically NOT `pattern = dsp`. DynamicSparsityPattern::operator= is documented to work
      // "only for empty objects" and enforces that with an Assert -- which is compiled out in
      // Release. So the obvious assignment silently leaves the pattern 0x0, PETSc then preallocates
      // from an uninitialised row-start array, and the failure surfaces as
      // "nnz cannot be greater than row length: value 545352472" from deep inside MatSeqAIJ.
      // Copy the entries explicitly instead.
      const auto &rows = dsp.row_index_set();
      pattern.reinit(dsp.n_rows(), dsp.n_cols(), rows);
      const auto copy_row = [&](const dealii::types::global_dof_index row) {
        for (auto it = dsp.begin(row); it != dsp.end(row); ++it)
          pattern.add(row, it->column());
      };
      if (rows.n_elements() > 0)
        for (const auto row : rows)
          copy_row(row);
      else
        for (dealii::types::global_dof_index row = 0; row < dsp.n_rows(); ++row)
          copy_row(row);
#endif
    } else {
      (void)locally_owned;
      (void)locally_relevant;
      (void)comm;
      pattern.copy_from(dsp);
    }
  }
 
  template <typename SparseMatrixType>
  void reinit_la_matrix(SparseMatrixType &matrix, const get_type::SparsityPattern<SparseMatrixType> &pattern,
                        const dealii::IndexSet &locally_owned, MPI_Comm comm)
  {
    if constexpr (is_distributed_la<SparseMatrixType>) {
      matrix.reinit(locally_owned, pattern, comm);
    } else {
      (void)locally_owned;
      (void)comm;
      matrix.reinit(pattern);
    }
  }
 
  template <typename VectorType>
  dealii::IndexSet restrict_to_owned(const dealii::IndexSet &global_set, const dealii::IndexSet &locally_owned)
  {
    if constexpr (is_distributed_la<VectorType>) {
      return global_set & locally_owned;
    } else {
      (void)locally_owned;
      return global_set;
    }
  }
 
  inline dealii::IndexSet variables_owner_set(const dealii::types::global_dof_index n_vars, MPI_Comm comm)
  {
    constexpr unsigned int variable_owner = 0;
    dealii::IndexSet variables(n_vars);
    if (MPI::rank(comm) == variable_owner) variables.add_range(0, n_vars);
    variables.compress();
    return variables;
  }
 
  template <typename VectorType>
  void reinit_variables_view(SolutionView<VectorType> &view, const dealii::types::global_dof_index n_vars,
                             MPI_Comm comm)
  {
    view.reinit(variables_owner_set(n_vars, comm), dealii::complete_index_set(n_vars), comm);
  }
 
  template <typename VectorType>
  void reinit_la_variables_vector(VectorType &vec, const dealii::types::global_dof_index n_vars, MPI_Comm comm)
  {
    if constexpr (is_distributed_la<VectorType>) {
      vec.reinit(variables_owner_set(n_vars, comm), comm);
    } else {
      (void)comm;
      vec.reinit(n_vars);
    }
  }
 
  template <typename BlockVectorType>
  void reinit_la_block_vector(BlockVectorType &vec, const std::vector<uint> &block_structure,
                              const dealii::IndexSet &locally_owned, MPI_Comm comm)
  {
    if constexpr (is_distributed_la<BlockVectorType>) {
      std::vector<dealii::IndexSet> owned{locally_owned};
      if (block_structure.size() > 1) owned.push_back(variables_owner_set(block_structure[1], comm));
      vec.reinit(owned, comm);
    } else {
      (void)locally_owned;
      (void)comm;
      vec.reinit(block_structure);
    }
  }
 
  template <typename VectorType, typename Fn>
  void compute_variables_into(VectorType &dst, VectorType &scratch, Fn &&compute)
  {
    if constexpr (!is_distributed_la<VectorType>) {
      (void)scratch;
      compute(dst);
    } else {
      scratch = 0;
      compute(scratch);
      // The model wrote through VectorBase::operator(), which stages rather than stores.
      scratch.compress(dealii::VectorOperation::insert);
      for (const auto i : dst.locally_owned_elements())
        dst(i) = scratch(i);
      dst.compress(dealii::VectorOperation::insert);
    }
  }
 
  template <typename VectorType>
  void reinit_local_variables_vector(VectorType &vec, const dealii::types::global_dof_index n_vars)
  {
    if constexpr (is_distributed_la<VectorType>) {
      if (n_vars > 0) vec.reinit(dealii::complete_index_set(n_vars), MPI_COMM_SELF);
    } else {
      vec.reinit(n_vars);
    }
  }
 
  template <typename VectorType, typename NumberType>
  void dense_vmult_variables(const dealii::FullMatrix<NumberType> &matrix, VectorType &dst, const VectorType &src,
                             MPI_Comm comm)
  {
    if constexpr (!is_distributed_la<VectorType>) {
      (void)comm;
      matrix.vmult(dst, src);
    } else {
      const auto n = src.size();
      dealii::Vector<NumberType> local_src(n), local_dst(n);
 
      const auto owned = src.locally_owned_elements();
      for (const auto i : owned)
        local_src[i] = src(i);
      MPI::bcast(comm, local_src.data(), static_cast<size_t>(n) * sizeof(NumberType), /*root=*/0);
 
      matrix.vmult(local_dst, local_src);
 
      for (const auto i : dst.locally_owned_elements())
        dst(i) = local_dst[i];
      dst.compress(dealii::VectorOperation::insert);
    }
  }
} // namespace DiFfRG