cuda_matrix_free.h

Go to the documentation of this file.

// ---------------------------------------------------------------------
//
// Copyright (C) 2016 - 2022 by the deal.II authors
//
// This file is part of the deal.II library.
//
// The deal.II library is free software; you can use it, redistribute
// it, and/or modify it under the terms of the GNU Lesser General
// Public License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
// The full text of the license can be found in the file LICENSE.md at
// the top level directory of deal.II.
//
// ---------------------------------------------------------------------
 
 
#ifndef dealii_cuda_matrix_free_h
#define dealii_cuda_matrix_free_h
 
#include <deal.II/base/config.h>
 
#ifdef DEAL_II_COMPILER_CUDA_AWARE
 
#  include <deal.II/base/cuda_size.h>
#  include <deal.II/base/mpi.h>
#  include <deal.II/base/quadrature.h>
#  include <deal.II/base/tensor.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/fe/mapping_q1.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>
 
 
 
DEAL_II_NAMESPACE_OPEN
 
// Forward declaration
namespace internal
{
  namespace MatrixFreeFunctions
  {
    enum class ConstraintKinds : std::uint16_t;
  }
} // namespace internal
 
namespace CUDAWrappers
{
  // 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>;
 
    enum ParallelizationScheme
    {
      parallel_in_elem,
      parallel_over_elem
    };
 
    struct AdditionalData
    {
      AdditionalData(
        const ParallelizationScheme parallelization_scheme = parallel_in_elem,
        const UpdateFlags           mapping_update_flags   = update_gradients |
                                                 update_JxW_values |
                                                 update_quadrature_points,
        const bool use_coloring                      = false,
        const bool overlap_communication_computation = false)
        : parallelization_scheme(parallelization_scheme)
        , mapping_update_flags(mapping_update_flags)
        , use_coloring(use_coloring)
        , overlap_communication_computation(overlap_communication_computation)
      {
#  ifndef DEAL_II_MPI_WITH_CUDA_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."));
      }
      ParallelizationScheme parallelization_scheme;
      UpdateFlags mapping_update_flags;
 
      bool use_coloring;
 
      bool overlap_communication_computation;
    };
 
    struct Data
    {
      point_type *q_points;
 
      types::global_dof_index *local_to_global;
 
      Number *inv_jacobian;
 
      Number *JxW;
 
      unsigned int id;
 
      unsigned int n_cells;
 
      unsigned int padding_length;
 
      unsigned int row_start;
 
      ::internal::MatrixFreeFunctions::ConstraintKinds *constraint_mask;
 
      bool use_coloring;
    };
 
    MatrixFree();
 
    ~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;
 
    void
    initialize_dof_vector(
      LinearAlgebra::CUDAWrappers::Vector<Number> &vec) const;
 
    void
    initialize_dof_vector(
      LinearAlgebra::distributed::Vector<Number, MemorySpace::CUDA> &vec) const;
 
    const std::vector<std::vector<CellFilter>> &
    get_colored_graph() const;
 
    const std::shared_ptr<const Utilities::MPI::Partitioner> &
    get_vector_partitioner() const;
 
    void
    free();
 
    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,
                    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::CUDA> &src,
      LinearAlgebra::distributed::Vector<Number, MemorySpace::CUDA> &dst) const;
 
    template <typename Functor>
    void
    distributed_cell_loop(
      const Functor &                                    func,
      const LinearAlgebra::CUDAWrappers::Vector<Number> &src,
      LinearAlgebra::CUDAWrappers::Vector<Number> &      dst) const;
 
    template <typename VectorType>
    void
    serial_copy_constrained_values(const VectorType &src,
                                   VectorType &      dst) const;
 
    void
    distributed_copy_constrained_values(
      const LinearAlgebra::distributed::Vector<Number, MemorySpace::CUDA> &src,
      LinearAlgebra::distributed::Vector<Number, MemorySpace::CUDA> &dst) const;
 
    void
    distributed_copy_constrained_values(
      const LinearAlgebra::CUDAWrappers::Vector<Number> &src,
      LinearAlgebra::CUDAWrappers::Vector<Number> &      dst) const;
 
    template <typename VectorType>
    void
    serial_set_constrained_values(const Number val, VectorType &dst) const;
 
    void
    distributed_set_constrained_values(
      const Number                                                   val,
      LinearAlgebra::distributed::Vector<Number, MemorySpace::CUDA> &dst) const;
 
    void
    distributed_set_constrained_values(
      const Number                                 val,
      LinearAlgebra::CUDAWrappers::Vector<Number> &dst) const;
 
    int my_id;
 
    ParallelizationScheme parallelization_scheme;
 
    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<point_type *> q_points;
 
    std::vector<types::global_dof_index *> local_to_global;
 
    std::vector<Number *> inv_jacobian;
 
    std::vector<Number *> JxW;
 
    types::global_dof_index *constrained_dofs;
 
    std::vector<::internal::MatrixFreeFunctions::ConstraintKinds *>
      constraint_mask;
 
    std::vector<dim3> grid_dim;
 
    std::vector<dim3> block_dim;
 
    std::shared_ptr<const Utilities::MPI::Partitioner> partitioner;
 
    unsigned int cells_per_block;
 
    dim3 constraint_grid_dim;
 
    dim3 constraint_block_dim;
 
    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>;
  };
 
 
 
  // TODO find a better place to put these things
 
  template <int dim, typename Number>
  struct SharedData
  {
    __device__
    SharedData(Number *vd, Number *gq[dim])
      : values(vd)
    {
      for (unsigned int d = 0; d < dim; ++d)
        gradients[d] = gq[d];
    }
 
    Number *values;
 
    Number *gradients[dim];
  };
 
 
 
  // This function determines the number of cells per block, possibly at compile
  // time (by virtue of being 'constexpr')
  // TODO this function should be rewritten using meta-programming
  __host__ __device__ constexpr unsigned int
  cells_per_block_shmem(int dim, int fe_degree)
  {
    /* clang-format off */
    // We are limiting the number of threads according to the
    // following formulas:
    //  - in 2D: `threads = cells * (k+1)^d <= 4*CUDAWrappers::warp_size`
    //  - in 3D: `threads = cells * (k+1)^d <= 2*CUDAWrappers::warp_size`
    return dim==2 ? (fe_degree==1 ? CUDAWrappers::warp_size :    // 128
                     fe_degree==2 ? CUDAWrappers::warp_size/4 :  //  72
                     fe_degree==3 ? CUDAWrappers::warp_size/8 :  //  64
                     fe_degree==4 ? CUDAWrappers::warp_size/8 :  // 100
                     1) :
           dim==3 ? (fe_degree==1 ? CUDAWrappers::warp_size/4 :  //  64
                     fe_degree==2 ? CUDAWrappers::warp_size/16 : //  54
                     1) : 1;
    /* clang-format on */
  }
 
 
  /*----------------------- Helper functions ---------------------------------*/
  template <int dim>
  __device__ inline unsigned int
  q_point_id_in_cell(const unsigned int n_q_points_1d)
  {
    return (
      dim == 1 ? threadIdx.x % n_q_points_1d :
      dim == 2 ? threadIdx.x % n_q_points_1d + n_q_points_1d * threadIdx.y :
                 threadIdx.x % n_q_points_1d +
                   n_q_points_1d * (threadIdx.y + n_q_points_1d * threadIdx.z));
  }
 
 
 
  template <int dim, typename Number>
  __device__ inline unsigned int
  local_q_point_id(
    const unsigned int                                          cell,
    const typename CUDAWrappers::MatrixFree<dim, Number>::Data *data,
    const unsigned int                                          n_q_points_1d,
    const unsigned int                                          n_q_points)
  {
    return (data->row_start / data->padding_length + cell) * n_q_points +
           q_point_id_in_cell<dim>(n_q_points_1d);
  }
 
 
 
  template <int dim, typename Number>
  __device__ inline typename CUDAWrappers::MatrixFree<dim, Number>::point_type &
  get_quadrature_point(
    const unsigned int                                          cell,
    const typename CUDAWrappers::MatrixFree<dim, Number>::Data *data,
    const unsigned int                                          n_q_points_1d)
  {
    return *(data->q_points + data->padding_length * cell +
             q_point_id_in_cell<dim>(n_q_points_1d));
  }
 
  template <int dim, typename Number>
  struct DataHost
  {
    std::vector<Point<dim, Number>> q_points;
 
    std::vector<types::global_dof_index> local_to_global;
 
    std::vector<Number> inv_jacobian;
 
    std::vector<Number> JxW;
 
    unsigned int id;
 
    unsigned int n_cells;
 
    unsigned int padding_length;
 
    unsigned int row_start;
 
    std::vector<::internal::MatrixFreeFunctions::ConstraintKinds>
      constraint_mask;
 
    bool use_coloring;
  };
 
 
 
  template <int dim, typename Number>
  DataHost<dim, Number>
  copy_mf_data_to_host(
    const typename ::CUDAWrappers::MatrixFree<dim, Number>::Data &data,
    const UpdateFlags &update_flags)
  {
    DataHost<dim, Number> data_host;
 
    data_host.id             = data.id;
    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;
 
    const unsigned int n_elements =
      data_host.n_cells * data_host.padding_length;
    if (update_flags & update_quadrature_points)
      {
        data_host.q_points.resize(n_elements);
        Utilities::CUDA::copy_to_host(data.q_points, data_host.q_points);
      }
 
    data_host.local_to_global.resize(n_elements);
    Utilities::CUDA::copy_to_host(data.local_to_global,
                                  data_host.local_to_global);
 
    if (update_flags & update_gradients)
      {
        data_host.inv_jacobian.resize(n_elements * dim * dim);
        Utilities::CUDA::copy_to_host(data.inv_jacobian,
                                      data_host.inv_jacobian);
      }
 
    if (update_flags & update_JxW_values)
      {
        data_host.JxW.resize(n_elements);
        Utilities::CUDA::copy_to_host(data.JxW, data_host.JxW);
      }
 
    data_host.constraint_mask.resize(data_host.n_cells);
    Utilities::CUDA::copy_to_host(data.constraint_mask,
                                  data_host.constraint_mask);
 
    return data_host;
  }
 
 
 
  template <int dim, typename Number>
  inline unsigned int
  local_q_point_id_host(const unsigned int           cell,
                        const DataHost<dim, Number> &data,
                        const unsigned int           n_q_points,
                        const unsigned int           i)
  {
    return (data.row_start / data.padding_length + cell) * n_q_points + i;
  }
 
 
 
  template <int dim, typename Number>
  inline Point<dim, Number>
  get_quadrature_point_host(const unsigned int           cell,
                            const DataHost<dim, Number> &data,
                            const unsigned int           i)
  {
    return data.q_points[data.padding_length * cell + i];
  }
 
 
  /*----------------------- 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 CUDAWrappers
 
DEAL_II_NAMESPACE_CLOSE
 
#endif
#endif