portable_matrix_free.h

Go to the documentation of this file.

// -----------------------------------------------------------------------------
//
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception OR LGPL-2.1-or-later
// Copyright (C) 2017 - 2026 by the deal.II authors
//
// This file is part of the deal.II library.
//
// Detailed license information governing the source code and contributions
// can be found in LICENSE.md and CONTRIBUTING.md at the top level directory.
//
// -----------------------------------------------------------------------------
 
 
#ifndef dealii_portable_matrix_free_h
#define dealii_portable_matrix_free_h
 
#include <deal.II/base/config.h>
 
#include <deal.II/base/exceptions.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/la_parallel_vector.h>
 
#include <deal.II/matrix_free/shape_info.h>
 
#include <Kokkos_Array.hpp>
#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;
  }
  template <int dim, typename Number>
  struct SharedData;
#endif

  inline constexpr unsigned int n_max_dof_handlers = 5;

  template <typename Number>
  using DeviceVector =
    Kokkos::View<Number *, MemorySpace::Default::kokkos_space>;

  template <typename Number>
  class DeviceBlockVector
  {
  public:
    DeviceBlockVector()
      : n_blocks(0)
    {}

    DeviceBlockVector(const DeviceBlockVector &other) = default;

    explicit DeviceBlockVector(const DeviceVector<Number> &src)
      : n_blocks(1)
      , blocks{src}
    {}

    template <typename MemorySpace>
    DeviceBlockVector(
      const LinearAlgebra::distributed::BlockVector<Number, MemorySpace> &src)
      : n_blocks(src.n_blocks())
    {
      Assert(src.n_blocks() <= n_max_dof_handlers,
             ExcMessage("Portable::MatrixFree is configured with " +
                        Utilities::to_string(n_max_dof_handlers) +
                        " but you are passing a BlockVector with " +
                        Utilities::to_string(src.n_blocks()) + " blocks."));
 
      for (int b = 0; b < src.n_blocks(); ++b)
        blocks[b] = DeviceVector<Number>(src.block(b).get_values(),
                                         src.block(b).locally_owned_size());
    }

    DEAL_II_HOST_DEVICE
    DeviceVector<Number> &
    block(unsigned int index)
    {
      AssertIndexRange(index, n_blocks);
      return blocks[index];
    }

    DEAL_II_HOST_DEVICE const DeviceVector<Number>                           &
    block(unsigned int index) const
    {
      AssertIndexRange(index, n_blocks);
      return blocks[index];
    }
 
  private:
    unsigned int n_blocks;

    Kokkos::Array<DeviceVector<Number>, n_max_dof_handlers> blocks;
  };

  template <int dim, typename Number = double>
  class MatrixFree : public EnableObserverPointer
  {
  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>;
    using LevelCellFilter =
      FilteredIterator<typename DoFHandler<dim>::level_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,
        const unsigned int mg_level  = numbers::invalid_unsigned_int,
        const unsigned int team_size = numbers::invalid_unsigned_int)
        : mapping_update_flags(mapping_update_flags)
        , use_coloring(use_coloring)
        , overlap_communication_computation(overlap_communication_computation)
        , mg_level(mg_level)
        , team_size(team_size)
      {}

      UpdateFlags mapping_update_flags;

      bool use_coloring;

      bool overlap_communication_computation;

      unsigned int mg_level;

      unsigned int team_size;
    };
 
 

    struct PrecomputedData
    {
      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 n_components;

      unsigned int padding_length;

      unsigned int row_start;

      bool use_coloring;

      ::internal::MatrixFreeFunctions::ElementType element_type;

      unsigned int scratch_pad_size;
    };
 

    struct Data
    {
      using TeamHandle = Kokkos::TeamPolicy<
        MemorySpace::Default::kokkos_space::execution_space>::member_type;

      TeamHandle team_member;
 
      const unsigned int n_q_points;
      const unsigned int n_dof_handler;
      const int          cell_index;
      const Kokkos::Array<PrecomputedData, n_max_dof_handlers>
        &precomputed_data;
      Kokkos::Array<SharedData<dim, Number>, n_max_dof_handlers> &shared_data;

      DEAL_II_DEPRECATED 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
      {
        Assert(n_q_points_ == n_q_points,
               ExcMessage("Incorrect argument value: n_q_points"));
        return local_q_point_id(cell, q_point);
      }

      DEAL_II_HOST_DEVICE unsigned int
      local_q_point_id(const unsigned int cell,
                       const unsigned int q_point) const
      {
        AssertIndexRange(cell, precomputed_data[0].n_cells);
        AssertIndexRange(q_point, n_q_points);
 
        Assert(precomputed_data[0].n_cells ==
                 precomputed_data[0].q_points.extent(1),
               ExcInternalError("q_points array has wrong size"));
        Assert(n_q_points == precomputed_data[0].q_points.extent(0),
               ExcInternalError("q_points array has wrong size"));
 
        return (precomputed_data[0].row_start /
                  precomputed_data[0].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
      {
        Assert(precomputed_data[0].n_cells ==
                 precomputed_data[0].q_points.extent(1),
               ExcInternalError());
        AssertIndexRange(cell, precomputed_data[0].n_cells);
        AssertIndexRange(q_point, n_q_points);
        Assert(n_q_points == precomputed_data[0].q_points.extent(0),
               ExcInternalError());
        return precomputed_data[0].q_points(q_point, cell);
      }

      template <typename Functor>
      DEAL_II_HOST_DEVICE void
      for_each_quad_point(const Functor &func) const
      {
        Kokkos::parallel_for(Kokkos::TeamThreadRange(team_member, n_q_points),
                             [&](const int &q) { func(q); });
        team_member.team_barrier();
      }
    };

    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());
 

    template <typename IteratorFiltersType>
    void
    reinit(const Mapping<dim>                                   &mapping,
           const std::vector<const DoFHandler<dim> *>           &dof_handler,
           const std::vector<const AffineConstraints<Number> *> &constraints,
           const Quadrature<1>                                  &quad,
           const IteratorFiltersType &iterator_filter,
           const AdditionalData      &additional_data = AdditionalData());

    void
    reinit(const Mapping<dim>                                   &mapping,
           const std::vector<const DoFHandler<dim> *>           &dof_handler,
           const std::vector<const AffineConstraints<Number> *> &constraints,
           const Quadrature<1>                                  &quad,
           const AdditionalData &additional_data = AdditionalData());

    PrecomputedData
    get_data(const unsigned int color,
             const unsigned int dof_handler_index = 0) 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 unsigned int dof_handler_index = 0) const;

    void
    copy_constrained_values(
      const LinearAlgebra::distributed::BlockVector<Number,
                                                    MemorySpace::Default> &src,
      LinearAlgebra::distributed::BlockVector<Number, MemorySpace::Default>
        &dst) const;

    template <typename VectorType>
    void
    set_constrained_values(const Number       value,
                           VectorType        &dst,
                           const unsigned int dof_handler_index = 0) const;

    template <typename MemorySpaceType>
    void
    initialize_dof_vector(
      LinearAlgebra::distributed::Vector<Number, MemorySpaceType> &vec,
      const unsigned int dof_handler_index = 0) const;

    template <typename MemorySpaceType>
    void
    initialize_dof_vector(
      LinearAlgebra::distributed::BlockVector<Number, MemorySpaceType> &vec)
      const;

    const std::vector<std::vector<CellFilter>> &
    get_colored_graph() const;

    const std::vector<std::vector<LevelCellFilter>> &
    get_colored_level_graph() const;

    const std::shared_ptr<const Utilities::MPI::Partitioner> &
    get_vector_partitioner(const unsigned int dof_handler_index = 0) const;

    const DoFHandler<dim> &
    get_dof_handler(const unsigned int dof_handler_index = 0) const;

    unsigned int
    get_mg_level() const;

    typename DoFHandler<dim>::cell_iterator
    get_cell_iterator(const unsigned int color,
                      const unsigned int index,
                      const unsigned int dof_handler_index = 0) const;

    unsigned int
    n_cells_per_color(const unsigned int color) const;

    bool
    use_overlap_communication_computation() const;

    std::size_t
    memory_consumption() const;
 
  private:
    template <typename IteratorFiltersType>
    void
    internal_reinit(
      const Mapping<dim>                                   &mapping,
      const std::vector<const DoFHandler<dim> *>           &dof_handler,
      const std::vector<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;

    template <typename Functor>
    void
    distributed_cell_loop(
      const Functor                                                       &func,
      const LinearAlgebra::distributed::BlockVector<Number,
                                                    MemorySpace::Default> &src,
      LinearAlgebra::distributed::BlockVector<Number, MemorySpace::Default>
        &dst) const;
 

    int my_id;

    bool use_coloring;

    bool overlap_communication_computation;

    unsigned int mg_level;

    unsigned int team_size;
 

    struct MappingInfo
    {
      std::vector<
        Kokkos::View<point_type **, MemorySpace::Default::kokkos_space>>
        q_points;

      std::vector<
        Kokkos::View<Number **[dim][dim], MemorySpace::Default::kokkos_space>>
        inv_jacobian;

      std::vector<Kokkos::View<Number **, MemorySpace::Default::kokkos_space>>
        JxW;
    };

    std::vector<MappingInfo> mapping_info;

    struct DoFInfo
    {
      types::global_dof_index n_dofs;

      ::internal::MatrixFreeFunctions::ElementType element_type;

      unsigned int scratch_pad_size;

      unsigned int fe_degree;

      unsigned int n_components;

      unsigned int scalar_dofs_per_cell;

      unsigned int dofs_per_cell;

      unsigned int n_constrained_dofs;

      std::vector<Kokkos::View<types::global_dof_index **,
                               MemorySpace::Default::kokkos_space>>
        local_to_global;

      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;

      const DoFHandler<dim> *dof_handler;

      std::shared_ptr<const Utilities::MPI::Partitioner> partitioner;
    };

    std::vector<DoFInfo> dof_handler_data;

    unsigned int q_points_per_cell;

    unsigned int n_colors;

    std::vector<unsigned int> n_cells;

    unsigned int padding_length;

    std::vector<unsigned int> row_start;
 

    std::vector<std::vector<CellFilter>> graph;

    std::vector<std::vector<LevelCellFilter>> level_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 SharedViewValues = Kokkos::View<
      Number **,
      MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space,
      Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
    using SharedViewGradients = Kokkos::View<
      Number ***,
      MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space,
      Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
    using SharedViewScratchPad = Kokkos::View<
      Number *,
      MemorySpace::Default::kokkos_space::execution_space::scratch_memory_space,
      Kokkos::MemoryTraits<Kokkos::Unmanaged>>;

    DEAL_II_HOST_DEVICE
    void
    reinit(const TeamHandle                                        &team_handle,
           const unsigned int                                       n_q_points,
           const typename MatrixFree<dim, Number>::PrecomputedData &gpu_data)
    {
      values      = SharedViewValues(team_handle.team_shmem(),
                                n_q_points,
                                gpu_data.n_components);
      gradients   = SharedViewGradients(team_handle.team_shmem(),
                                      n_q_points,
                                      dim,
                                      gpu_data.n_components);
      scratch_pad = SharedViewScratchPad(team_handle.team_shmem(),
                                         gpu_data.scratch_pad_size);
    }

    SharedViewValues values;

    SharedViewGradients gradients;

    SharedViewScratchPad scratch_pad;
  };
 
 

  template <int dim, typename Number>
  struct DataHost
  {
    typename Kokkos::View<Point<dim, Number> **,
                          MemorySpace::Default::kokkos_space>::host_mirror_type
      q_points;

    typename Kokkos::View<types::global_dof_index **,
                          MemorySpace::Default::kokkos_space>::host_mirror_type
      local_to_global;

    typename Kokkos::View<Number **[dim][dim],
                          MemorySpace::Default::kokkos_space>::host_mirror_type
      inv_jacobian;

    typename Kokkos::View<Number **,
                          MemorySpace::Default::kokkos_space>::host_mirror_type
      JxW;

    unsigned int n_cells;

    unsigned int padding_length;

    unsigned int row_start;

    typename Kokkos::View<
      ::internal::MatrixFreeFunctions::ConstraintKinds *,
      MemorySpace::Default::kokkos_space>::host_mirror_type 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(q_point, cell);
    }
  };
 
 

  template <int dim, typename Number>
  DataHost<dim, Number>
  copy_mf_data_to_host(
    const typename ::Portable::MatrixFree<dim, Number>::PrecomputedData
                      &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
  {
    AssertThrow(
      mg_level == numbers::invalid_unsigned_int,
      ExcMessage(
        "The MatrixFree object has been initialized for a specific "
        "multigrid level. This function is only available when working "
        "on active cells."));
    return graph;
  }
 
 
 
  template <int dim, typename Number>
  inline const std::vector<std::vector<
    FilteredIterator<typename DoFHandler<dim>::level_cell_iterator>>> &
  MatrixFree<dim, Number>::get_colored_level_graph() const
  {
    AssertThrow(
      mg_level != numbers::invalid_unsigned_int,
      ExcMessage("The MatrixFree object has been initialized for active cells. "
                 "This function is only available when working on a specific "
                 "multigrid level."));
    return level_graph;
  }
 
 
 
  template <int dim, typename Number>
  inline const std::shared_ptr<const Utilities::MPI::Partitioner> &
  MatrixFree<dim, Number>::get_vector_partitioner(
    const unsigned int dof_handler_index) const
  {
    Assert(dof_handler_data.size() > 0, ExcNotInitialized());
    AssertIndexRange(dof_handler_index, dof_handler_data.size());
    return dof_handler_data[dof_handler_index].partitioner;
  }
 
 
 
  template <int dim, typename Number>
  inline unsigned int
  MatrixFree<dim, Number>::get_mg_level() const
  {
    return mg_level;
  }
 
 
 
  template <int dim, typename Number>
  inline const DoFHandler<dim> &
  MatrixFree<dim, Number>::get_dof_handler(
    const unsigned int dof_handler_index) const
  {
    Assert(dof_handler_data.size() > 0, ExcNotInitialized());
    AssertIndexRange(dof_handler_index, dof_handler_data.size());
 
    return *dof_handler_data[dof_handler_index].dof_handler;
  }
 
 
  template <int dim, typename Number>
  typename DoFHandler<dim>::cell_iterator
  MatrixFree<dim, Number>::get_cell_iterator(
    const unsigned int color,
    const unsigned int index,
    const unsigned int dof_handler_index) const
  {
    if (mg_level == numbers::invalid_unsigned_int)
      {
        AssertIndexRange(color, get_colored_graph().size());
        AssertIndexRange(index, get_colored_graph()[color].size());
 
        const unsigned int cell_level =
          get_colored_graph()[color][index]->level();
 
        const unsigned int cell_index =
          get_colored_graph()[color][index]->index();
 
        return typename DoFHandler<dim>::cell_iterator(
          &(get_dof_handler(dof_handler_index).get_triangulation()),
          cell_level,
          cell_index,
          &get_dof_handler(dof_handler_index));
      }
    else
      {
        AssertIndexRange(color, get_colored_level_graph().size());
        AssertIndexRange(index, get_colored_level_graph()[color].size());
 
        const unsigned int cell_level =
          get_colored_level_graph()[color][index]->level();
 
        const unsigned int cell_index =
          get_colored_level_graph()[color][index]->index();
 
        return typename DoFHandler<dim>::cell_iterator(
          &(get_dof_handler(dof_handler_index).get_triangulation()),
          cell_level,
          cell_index,
          &get_dof_handler(dof_handler_index));
      }
  }
 
 
 
  template <int dim, typename Number>
  inline unsigned int
  MatrixFree<dim, Number>::n_cells_per_color(const unsigned int color) const
  {
    AssertIndexRange(color, n_cells.size());
    return n_cells[color];
  }
 
 
 
  template <int dim, typename Number>
  inline bool
  MatrixFree<dim, Number>::use_overlap_communication_computation() const
  {
    return overlap_communication_computation;
  }
 
#endif
 
} // namespace Portable
 
DEAL_II_NAMESPACE_CLOSE
 
#endif