cuda_matrix_free.h

// ---------------------------------------------------------------------
//
// Copyright (C) 2016 - 2018 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 at
// the top level of the deal.II distribution.
//
// ---------------------------------------------------------------------


#ifndef dealii_cuda_matrix_free_h
#define dealii_cuda_matrix_free_h

#include <deal.II/base/config.h>

#ifdef DEAL_II_WITH_CUDA

#include <deal.II/base/quadrature.h>
#include <deal.II/base/tensor.h>
#include <deal.II/dofs/dof_handler.h>
#include <deal.II/fe/mapping.h>
#include <deal.II/fe/mapping_q1.h>
#include <deal.II/fe/fe_update_flags.h>
#include <deal.II/lac/constraint_matrix.h>
#include <deal.II/lac/cuda_vector.h>
#include <cuda_runtime_api.h>


DEAL_II_NAMESPACE_OPEN

namespace CUDAWrappers
{
  // forward declaration
  namespace internal
  {
    template <int dim, typename Number>
    class ReinitHelper;
  }

  template <int dim, typename Number=double>
  class MatrixFree : public Subscriptor
  {
  public:
    typedef Tensor<2, dim, Tensor<1,dim,Number>> jacobian_type;
    // TODO this should really be a CUDAWrappers::Point
    typedef Tensor<1, dim, Number> point_type;

    // Use Number2 so we don't hide the template parameter Number
    template <typename Number2>
    using CUDAVector = ::LinearAlgebra::CUDAWrappers::Vector<Number2>;

    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)
        :
        parallelization_scheme(parallelization_scheme),
        mapping_update_flags(mapping_update_flags)
      {}

      unsigned int n_colors;
      ParallelizationScheme parallelization_scheme;
      UpdateFlags mapping_update_flags;
    };

    struct Data
    {
      point_type *q_points;
      unsigned int *local_to_global;
      Number *inv_jacobian;
      Number *JxW;
      unsigned int n_cells;
      unsigned int padding_length;
      unsigned int row_start;
      unsigned int *constraint_mask;
    };

    MatrixFree();

    unsigned int get_padding_length() const;

    void reinit(const Mapping<dim> &mapping,
                const DoFHandler<dim> &dof_handler,
                const ConstraintMatrix &constraints,
                const Quadrature<1> &quad,
                const AdditionalData additional_data = AdditionalData());

    void reinit(const DoFHandler<dim> &dof_handler,
                const ConstraintMatrix &constraints,
                const Quadrature<1> &quad,
                const AdditionalData AdditionalData = AdditionalData());

    Data get_data(unsigned int color) const;

    template <typename functor>
    void cell_loop(const functor &func,
                   const CUDAVector<Number> &src,
                   CUDAVector<Number> &dst) const;

    void copy_constrained_values(const CUDAVector<Number> &src,
                                 CUDAVector<Number> &dst) const;

    void set_constrained_values(const Number val, CUDAVector<Number> &dst) const;

    void free();

    std::size_t memory_consumption() const;

  private:
    ParallelizationScheme parallelization_scheme;
    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<unsigned int *> local_to_global;
    std::vector<Number *> inv_jacobian;
    std::vector<Number *> JxW;

    // Constraints
    unsigned int *constrained_dofs;
    std::vector<unsigned int *> constraint_mask;
    std::vector<dim3> grid_dim;
    std::vector<dim3> block_dim;

    // Parallelization parameter
    unsigned int cells_per_block;
    dim3 constraint_grid_dim;
    dim3 constraint_block_dim;

    unsigned int padding_length;
    std::vector<unsigned int> row_start;

    friend class internal::ReinitHelper<dim,Number>;
  };



  // TODO find a better place to put these things
  // Structure to pass the shared memory into a general user function.
  template <int dim, typename Number>
  struct SharedData
  {
    __device__ SharedData(Number *vd,
                          Number *gq[dim])
      :
      values(vd)
    {
      for (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
  // TODO this function should be rewritten using meta-programming
  __host__ __device__ constexpr unsigned int cells_per_block_shmem(int dim,
      int fe_degree)
  {
    return dim==2 ? (fe_degree==1 ? 32 :
                     fe_degree==2 ? 8 :
                     fe_degree==3 ? 4 :
                     fe_degree==4 ? 4 :
                     1) :
           dim==3 ? (fe_degree==1 ? 8 :
                     fe_degree==2 ? 2 :
                     1) : 1;
  }
}

DEAL_II_NAMESPACE_CLOSE

#endif

#endif