portable_matrix_free.h

Go to the documentation of this file.

// ------------------------------------------------------------------------
//
// SPDX-License-Identifier: LGPL-2.1-or-later
// Copyright (C) 2017 - 2024 by the deal.II authors
//
// This file is part of the deal.II library.
//
// Part of the source code is dual licensed under Apache-2.0 WITH
// LLVM-exception OR LGPL-2.1-or-later. Detailed license information
// governing the source code and code contributions can be found in
// LICENSE.md and CONTRIBUTING.md at the top level directory of deal.II.
//
// ------------------------------------------------------------------------
 
 
#ifndef dealii_portable_matrix_free_h
#define dealii_portable_matrix_free_h
 
#include <deal.II/base/config.h>
 
#include <deal.II/base/cuda_size.h>
#include <deal.II/base/memory_space.h>
#include <deal.II/base/mpi_stub.h>
#include <deal.II/base/partitioner.h>
#include <deal.II/base/quadrature.h>
#include <deal.II/base/tensor.h>
#include <deal.II/base/utilities.h>
 
#include <deal.II/dofs/dof_handler.h>
 
#include <deal.II/fe/fe_update_flags.h>
#include <deal.II/fe/mapping.h>
 
#include <deal.II/grid/filtered_iterator.h>
 
#include <deal.II/lac/affine_constraints.h>
#include <deal.II/lac/cuda_vector.h>
#include <deal.II/lac/la_parallel_vector.h>
 
#include <Kokkos_Core.hpp>
 
 
 
DEAL_II_NAMESPACE_OPEN
 
// Forward declaration
namespace internal
{
  namespace MatrixFreeFunctions
  {
    enum class ConstraintKinds : std::uint16_t;
  }
} // namespace internal
 
namespace Portable
{
  // forward declaration
#ifndef DOXYGEN
  namespace internal
  {
    template <int dim, typename Number>
    class ReinitHelper;
  }
#endif
 
  template <int dim, typename Number = double>
  class MatrixFree : public Subscriptor
  {
  public:
    using jacobian_type = Tensor<2, dim, Tensor<1, dim, Number>>;
    using point_type    = Point<dim, Number>;
    using CellFilter =
      FilteredIterator<typename DoFHandler<dim>::active_cell_iterator>;
 
    struct AdditionalData
    {
      AdditionalData(const UpdateFlags mapping_update_flags =
                       update_gradients | update_JxW_values |
                       update_quadrature_points,
                     const bool use_coloring                      = false,
                     const bool overlap_communication_computation = false)
        : mapping_update_flags(mapping_update_flags)
        , use_coloring(use_coloring)
        , overlap_communication_computation(overlap_communication_computation)
      {
#ifndef DEAL_II_MPI_WITH_DEVICE_SUPPORT
        AssertThrow(
          overlap_communication_computation == false,
          ExcMessage(
            "Overlapping communication and computation requires CUDA-aware MPI."));
#endif
        if (overlap_communication_computation == true)
          AssertThrow(
            use_coloring == false || overlap_communication_computation == false,
            ExcMessage(
              "Overlapping communication and coloring are incompatible options. Only one of them can be enabled."));
      }
      UpdateFlags mapping_update_flags;
 
      bool use_coloring;
 
      bool overlap_communication_computation;
    };
 
    struct Data
    {
      Kokkos::View<point_type **, MemorySpace::Default::kokkos_space> q_points;
 
      Kokkos::View<types::global_dof_index **,
                   MemorySpace::Default::kokkos_space>
        local_to_global;
 
      Kokkos::View<Number **[dim][dim], MemorySpace::Default::kokkos_space>
        inv_jacobian;
 
      Kokkos::View<Number **, MemorySpace::Default::kokkos_space> JxW;
 
      Kokkos::View<::internal::MatrixFreeFunctions::ConstraintKinds *,
                   MemorySpace::Default::kokkos_space>
        constraint_mask;
 
      Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_values;
 
      Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
        shape_gradients;
 
      Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
        co_shape_gradients;
 
      Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
        constraint_weights;
 
      unsigned int n_cells;
 
      unsigned int padding_length;
 
      unsigned int row_start;
 
      bool use_coloring;
 
      DEAL_II_HOST_DEVICE unsigned int
      local_q_point_id(const unsigned int cell,
                       const unsigned int n_q_points,
                       const unsigned int q_point) const
      {
        return (row_start / padding_length + cell) * n_q_points + q_point;
      }
 
 
      DEAL_II_HOST_DEVICE
      typename Portable::MatrixFree<dim, Number>::point_type &
      get_quadrature_point(const unsigned int cell,
                           const unsigned int q_point) const
      {
        return q_points(cell, q_point);
      }
    };
 
    MatrixFree();
 
    unsigned int
    get_padding_length() const;
 
    template <typename IteratorFiltersType>
    void
    reinit(const Mapping<dim>              &mapping,
           const DoFHandler<dim>           &dof_handler,
           const AffineConstraints<Number> &constraints,
           const Quadrature<1>             &quad,
           const IteratorFiltersType       &iterator_filter,
           const AdditionalData            &additional_data = AdditionalData());
 
    void
    reinit(const Mapping<dim>              &mapping,
           const DoFHandler<dim>           &dof_handler,
           const AffineConstraints<Number> &constraints,
           const Quadrature<1>             &quad,
           const AdditionalData            &additional_data = AdditionalData());
 
    void
    reinit(const DoFHandler<dim>           &dof_handler,
           const AffineConstraints<Number> &constraints,
           const Quadrature<1>             &quad,
           const AdditionalData            &additional_data = AdditionalData());
 
    Data
    get_data(unsigned int color) const;
 
    // clang-format off
    // clang-format on
    template <typename Functor, typename VectorType>
    void
    cell_loop(const Functor    &func,
              const VectorType &src,
              VectorType       &dst) const;
 
    template <typename Functor>
    void
    evaluate_coefficients(Functor func) const;
 
    template <typename VectorType>
    void
    copy_constrained_values(const VectorType &src, VectorType &dst) const;
 
    template <typename VectorType>
    void
    set_constrained_values(const Number val, VectorType &dst) const;
 
#ifdef DEAL_II_WITH_CUDA
    void
    initialize_dof_vector(
      LinearAlgebra::CUDAWrappers::Vector<Number> &vec) const;
#endif
 
    void
    initialize_dof_vector(
      LinearAlgebra::distributed::Vector<Number, MemorySpace::Default> &vec)
      const;
 
    const std::vector<std::vector<CellFilter>> &
    get_colored_graph() const;
 
    const std::shared_ptr<const Utilities::MPI::Partitioner> &
    get_vector_partitioner() const;
 
    const DoFHandler<dim> &
    get_dof_handler() const;
 
    std::size_t
    memory_consumption() const;
 
  private:
    template <typename IteratorFiltersType>
    void
    internal_reinit(const Mapping<dim>                    &mapping,
                    const DoFHandler<dim>                 &dof_handler,
                    const AffineConstraints<Number>       &constraints,
                    const Quadrature<1>                   &quad,
                    const IteratorFiltersType             &iterator_filter,
                    const std::shared_ptr<const MPI_Comm> &comm,
                    const AdditionalData                   additional_data);
 
    template <typename Functor, typename VectorType>
    void
    serial_cell_loop(const Functor    &func,
                     const VectorType &src,
                     VectorType       &dst) const;
 
    template <typename Functor>
    void
    distributed_cell_loop(
      const Functor &func,
      const LinearAlgebra::distributed::Vector<Number, MemorySpace::Default>
                                                                       &src,
      LinearAlgebra::distributed::Vector<Number, MemorySpace::Default> &dst)
      const;
 
#ifdef DEAL_II_WITH_CUDA
    template <typename Functor>
    void
    distributed_cell_loop(
      const Functor                                     &func,
      const LinearAlgebra::CUDAWrappers::Vector<Number> &src,
      LinearAlgebra::CUDAWrappers::Vector<Number>       &dst) const;
#endif
 
    int my_id;
 
    bool use_coloring;
 
    bool overlap_communication_computation;
 
    types::global_dof_index n_dofs;
 
    unsigned int fe_degree;
 
    unsigned int dofs_per_cell;
 
    unsigned int n_constrained_dofs;
 
    unsigned int q_points_per_cell;
 
    unsigned int n_colors;
 
    std::vector<unsigned int> n_cells;
 
    std::vector<Kokkos::View<point_type **, MemorySpace::Default::kokkos_space>>
      q_points;
 
    std::vector<Kokkos::View<types::global_dof_index **,
                             MemorySpace::Default::kokkos_space>>
      local_to_global;
 
    std::vector<
      Kokkos::View<Number **[dim][dim], MemorySpace::Default::kokkos_space>>
      inv_jacobian;
 
    std::vector<Kokkos::View<Number **, MemorySpace::Default::kokkos_space>>
      JxW;
 
    Kokkos::View<types::global_dof_index *, MemorySpace::Default::kokkos_space>
      constrained_dofs;
 
    std::vector<
      Kokkos::View<::internal::MatrixFreeFunctions::ConstraintKinds *,
                   MemorySpace::Default::kokkos_space>>
      constraint_mask;
 
    Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_values;
 
    Kokkos::View<Number *, MemorySpace::Default::kokkos_space> shape_gradients;
 
    Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
      co_shape_gradients;
 
    Kokkos::View<Number *, MemorySpace::Default::kokkos_space>
      constraint_weights;
 
    std::shared_ptr<const Utilities::MPI::Partitioner> partitioner;
 
 
    unsigned int padding_length;
 
    std::vector<unsigned int> row_start;
 
    const DoFHandler<dim> *dof_handler;
 
    std::vector<std::vector<CellFilter>> graph;
 
    friend class internal::ReinitHelper<dim, Number>;
  };
 
 
 
  template <int dim, typename Number>
  struct SharedData
  {
    using TeamHandle = Kokkos::TeamPolicy<
      MemorySpace::Default::kokkos_space::execution_space>::member_type;
 
    using SharedView1D = Kokkos::View<
      Number *,
      MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space,
      Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
    using SharedView2D = Kokkos::View<
      Number *[dim],
      MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space,
      Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
 
    DEAL_II_HOST_DEVICE
    SharedData(const TeamHandle   &team_member,
               const SharedView1D &values,
               const SharedView2D &gradients)
      : team_member(team_member)
      , values(values)
      , gradients(gradients)
    {}
 
    TeamHandle team_member;
 
    SharedView1D values;
 
    SharedView2D gradients;
  };
 
 
 
  template <int dim, typename Number>
  struct DataHost
  {
    typename Kokkos::View<Point<dim, Number> **,
                          MemorySpace::Default::kokkos_space>::HostMirror
      q_points;
 
    typename Kokkos::View<types::global_dof_index **,
                          MemorySpace::Default::kokkos_space>::HostMirror
      local_to_global;
 
    typename Kokkos::View<Number **[dim][dim],
                          MemorySpace::Default::kokkos_space>::HostMirror
      inv_jacobian;
 
    typename Kokkos::View<Number **,
                          MemorySpace::Default::kokkos_space>::HostMirror JxW;
 
    unsigned int n_cells;
 
    unsigned int padding_length;
 
    unsigned int row_start;
 
    typename Kokkos::View<
      ::internal::MatrixFreeFunctions::ConstraintKinds *,
      MemorySpace::Default::kokkos_space>::HostMirror constraint_mask;
 
    bool use_coloring;
 
 
 
    unsigned int
    local_q_point_id(const unsigned int cell,
                     const unsigned int n_q_points,
                     const unsigned int q_point) const
    {
      return (row_start / padding_length + cell) * n_q_points + q_point;
    }
 
 
 
    Point<dim, Number>
    get_quadrature_point(const unsigned int cell,
                         const unsigned int q_point) const
    {
      return q_points(cell, q_point);
    }
  };
 
 
 
  template <int dim, typename Number>
  DataHost<dim, Number>
  copy_mf_data_to_host(
    const typename ::Portable::MatrixFree<dim, Number>::Data &data,
    const UpdateFlags &update_flags)
  {
    DataHost<dim, Number> data_host;
 
    data_host.n_cells        = data.n_cells;
    data_host.padding_length = data.padding_length;
    data_host.row_start      = data.row_start;
    data_host.use_coloring   = data.use_coloring;
 
    if (update_flags & update_quadrature_points)
      {
        data_host.q_points = Kokkos::create_mirror(data.q_points);
        Kokkos::deep_copy(data_host.q_points, data.q_points);
      }
 
    data_host.local_to_global = Kokkos::create_mirror(data.local_to_global);
    Kokkos::deep_copy(data_host.local_to_global, data.local_to_global);
 
    if (update_flags & update_gradients)
      {
        data_host.inv_jacobian = Kokkos::create_mirror(data.inv_jacobian);
        Kokkos::deep_copy(data_host.inv_jacobian, data.inv_jacobian);
      }
 
    if (update_flags & update_JxW_values)
      {
        data_host.JxW = Kokkos::create_mirror(data.JxW);
        Kokkos::deep_copy(data_host.JxW, data.JxW);
      }
 
    data_host.constraint_mask = Kokkos::create_mirror(data.constraint_mask);
    Kokkos::deep_copy(data_host.constraint_mask, data.constraint_mask);
 
    return data_host;
  }
 
 
  /*----------------------- Inline functions ---------------------------------*/
 
#ifndef DOXYGEN
 
  template <int dim, typename Number>
  inline const std::vector<std::vector<
    FilteredIterator<typename DoFHandler<dim>::active_cell_iterator>>> &
  MatrixFree<dim, Number>::get_colored_graph() const
  {
    return graph;
  }
 
 
 
  template <int dim, typename Number>
  inline const std::shared_ptr<const Utilities::MPI::Partitioner> &
  MatrixFree<dim, Number>::get_vector_partitioner() const
  {
    return partitioner;
  }
 
 
 
  template <int dim, typename Number>
  inline const DoFHandler<dim> &
  MatrixFree<dim, Number>::get_dof_handler() const
  {
    Assert(dof_handler != nullptr, ExcNotInitialized());
 
    return *dof_handler;
  }
 
#endif
 
} // namespace Portable
 
DEAL_II_NAMESPACE_CLOSE
 
#endif