Add MGTwoLevelTransferCopyToHost, a two level transfer class which permits transfer on device vectors via use of an internal MGTwoLevelTransfer class defined on the host.

Author: RyanMoulday

include/deal.II/multigrid/mg_transfer_matrix_free.h

@@ -133,26 +133,26 @@ class MGTwoLevelTransferBase : public EnableObserverPointer
     "type LinearAlgebra::distributed::Vector.");
 
   /**
-   * The scalar type used by the vector-type template argument.
+   * Partitioner needed by the intermediate vector.
    */
-  using Number = typename VectorType::value_type;
+  std::shared_ptr<const Utilities::MPI::Partitioner> partitioner_coarse;
 
   /**
-   * Default constructor.
+   * Partitioner needed by the intermediate vector.
    */
-  MGTwoLevelTransferBase();
+  std::shared_ptr<const Utilities::MPI::Partitioner> partitioner_fine;
 
   /**
    * Perform prolongation on a solution vector.
    */
-  void
-  prolongate_and_add(VectorType &dst, const VectorType &src) const;
+  virtual void
+  prolongate_and_add(VectorType &dst, const VectorType &src) const = 0;
 
   /**
    * Perform restriction on a residual vector.
    */
-  void
-  restrict_and_add(VectorType &dst, const VectorType &src) const;
+  virtual void
+  restrict_and_add(VectorType &dst, const VectorType &src) const = 0;
 
   /**
    * Perform interpolation of a solution vector from the fine level to the
@@ -181,6 +181,38 @@ class MGTwoLevelTransferBase : public EnableObserverPointer
    */
   virtual std::size_t
   memory_consumption() const = 0;
+};
+
+
+
+/**
+ * Base class for MGTwoLevelTransferNonNested and MGTwoLevelTransfer.
+ */
+template <typename VectorType>
+class MGTwoLevelTransferCore : public MGTwoLevelTransferBase<VectorType>
+{
+public:
+  /**
+   * The scalar type used by the vector-type template argument.
+   */
+  using Number = typename VectorType::value_type;
+
+  /**
+   * Default constructor.
+   */
+  MGTwoLevelTransferCore();
+
+  /**
+   * @copydoc MGTwoLevelTransferBase::prolongate_and_add
+   */
+  void
+  prolongate_and_add(VectorType &dst, const VectorType &src) const override;
+
+  /**
+   * @copydoc MGTwoLevelTransferBase::restrict_and_add
+   */
+  void
+  restrict_and_add(VectorType &dst, const VectorType &src) const override;
 
 protected:
   /**
@@ -242,17 +274,6 @@ class MGTwoLevelTransferBase : public EnableObserverPointer
    */
   bool fine_element_is_continuous;
 
-public:
-  /**
-   * Partitioner needed by the intermediate vector.
-   */
-  std::shared_ptr<const Utilities::MPI::Partitioner> partitioner_coarse;
-
-  /**
-   * Partitioner needed by the intermediate vector.
-   */
-  std::shared_ptr<const Utilities::MPI::Partitioner> partitioner_fine;
-
 protected:
   /**
    * Internal vector on which the actual prolongation/restriction is performed.
@@ -327,7 +348,7 @@ class MGTwoLevelTransferBase : public EnableObserverPointer
  * point, and we fall back to the first option in such a case.
  */
 template <int dim, typename VectorType>
-class MGTwoLevelTransfer : public MGTwoLevelTransferBase<VectorType>
+class MGTwoLevelTransfer : public MGTwoLevelTransferCore<VectorType>
 {
 public:
   static_assert(
@@ -627,6 +648,133 @@ class MGTwoLevelTransfer : public MGTwoLevelTransferBase<VectorType>
 
   friend class internal::MGTwoLevelTransferImplementation;
 
+  friend class MGTransferMatrixFree<dim, Number, MemorySpace::Host>;
+
+  friend class MGTransferMatrixFree<dim, Number, MemorySpace::Default>;
+};
+
+
+/**
+ * A transfer class supporting device vectors via copying them to the host and
+ * using an internal transfer of type MGTwoLevelTransfer defined on the
+ * host.
+ */
+template <int dim, typename VectorType>
+class MGTwoLevelTransferCopyToHost : public MGTwoLevelTransferBase<VectorType>
+{
+public:
+  static_assert(
+    std::is_same_v<
+      VectorType,
+      LinearAlgebra::distributed::Vector<typename VectorType::value_type,
+                                         MemorySpace::Default>>,
+    "This class is currently only implemented for vectors of "
+    "type LinearAlgebra::distributed::Vector, This class should not be initialized on the host as it will perform redundant transfers. Use MGTwoLevelTransfer instead.");
+
+
+  /**
+   * The scalar type used by the vector-type template argument.
+   */
+  using Number = typename VectorType::value_type;
+
+  using VectorTypeHost =
+    LinearAlgebra::distributed::Vector<Number, MemorySpace::Host>;
+  /**
+   * Default constructor.
+   */
+  MGTwoLevelTransferCopyToHost() = default;
+
+  /**
+   * @copydoc MGTwoLevelTransfer::reinit(const DoFHandler<dim>, const DoFHandler<dim>, const AffineConstraints<Number>, const AffineConstraints<Number>, const unsigned int, const unsigned int)
+   */
+  void
+  reinit(const DoFHandler<dim>           &dof_handler_fine,
+         const DoFHandler<dim>           &dof_handler_coarse,
+         const AffineConstraints<Number> &constraint_fine =
+           AffineConstraints<Number>(),
+         const AffineConstraints<Number> &constraint_coarse =
+           AffineConstraints<Number>(),
+         const unsigned int mg_level_fine   = numbers::invalid_unsigned_int,
+         const unsigned int mg_level_coarse = numbers::invalid_unsigned_int);
+
+  /**
+   * @copydoc MGTwoLevelTransfer::reinit(const MatrixFree<dim, Number>, const unsigned int, const MatrixFree<dim, Number>, const unsigned int)
+   */
+  void
+  reinit(const MatrixFree<dim, Number> &matrix_free_fine,
+         const unsigned int             dof_no_fine,
+         const MatrixFree<dim, Number> &matrix_free_coarse,
+         const unsigned int             dof_no_coarse);
+
+  /**
+   * @copydoc MGTwoLevelTransfer::fast_polynomial_transfer_supported
+   */
+  bool
+  fast_polynomial_transfer_supported(const unsigned int fe_degree_fine,
+                                     const unsigned int fe_degree_coarse);
+
+  /**
+   * @copydoc MGTwoLevelTransfer::prolongate_and_add
+   */
+  void
+  prolongate_and_add(VectorType &dst, const VectorType &src) const override;
+
+  /**
+   * @copydoc MGTwoLevelTransfer::restrict_and_add
+   */
+  void
+  restrict_and_add(VectorType &dst, const VectorType &src) const override;
+
+  /**
+   * @copydoc MGTwoLevelTransfer::interpolate
+   */
+  void
+  interpolate(VectorType &dst, const VectorType &src) const override;
+
+  /**
+   * @copydoc MGTwoLevelTransfer::enable_inplace_operations_if_possible
+   */
+  std::pair<bool, bool>
+  enable_inplace_operations_if_possible(
+    const std::shared_ptr<const Utilities::MPI::Partitioner>
+      &partitioner_coarse,
+    const std::shared_ptr<const Utilities::MPI::Partitioner> &partitioner_fine)
+    override;
+
+  /**
+   * @copydoc MGTwoLevelTransfer::memory_consumption
+   */
+  std::size_t
+  memory_consumption() const override;
+
+private:
+  /**
+   * The internal transfer defined on the host which handles all operations.
+   */
+  MGTwoLevelTransfer<dim, VectorTypeHost> host_transfer;
+
+  /**
+   * Vectors used for the host transfer.
+   */
+
+  mutable VectorTypeHost host_vector_coarse;
+
+  mutable VectorTypeHost host_vector_fine;
+
+  /**
+   * Copies a device vector to a host vector.
+   */
+  void
+  copy_to_host(VectorTypeHost &dst, const VectorType &src) const;
+
+  /**
+   * Copies a host vector to a device vector.
+   */
+  void
+  copy_from_host(VectorType &dst, const VectorTypeHost &src) const;
+
+  friend class internal::MGTwoLevelTransferImplementation;
+
   friend class MGTransferMatrixFree<dim,
                                     Number,
                                     typename VectorType::memory_space>;
@@ -638,7 +786,7 @@ class MGTwoLevelTransfer : public MGTwoLevelTransferBase<VectorType>
  * Class for transfer between two non-nested multigrid levels.
  */
 template <int dim, typename VectorType>
-class MGTwoLevelTransferNonNested : public MGTwoLevelTransferBase<VectorType>
+class MGTwoLevelTransferNonNested : public MGTwoLevelTransferCore<VectorType>
 {
 private:
   static_assert(
@@ -1451,6 +1599,150 @@ class MGTransferBlockMatrixFree
 
 /* ----------------------- Inline functions --------------------------------- */
 
+template <int dim, typename VectorType>
+void
+MGTwoLevelTransferCopyToHost<dim, VectorType>::reinit(
+  const DoFHandler<dim>           &dof_handler_fine,
+  const DoFHandler<dim>           &dof_handler_coarse,
+  const AffineConstraints<Number> &constraint_fine,
+  const AffineConstraints<Number> &constraint_coarse,
+  const unsigned int               mg_level_fine,
+  const unsigned int               mg_level_coarse)
+{
+  host_transfer.reinit(dof_handler_fine,
+                       dof_handler_coarse,
+                       constraint_fine,
+                       constraint_coarse,
+                       mg_level_fine,
+                       mg_level_coarse);
+}
+
+template <int dim, typename VectorType>
+void
+MGTwoLevelTransferCopyToHost<dim, VectorType>::reinit(
+  const MatrixFree<dim, Number> &matrix_free_fine,
+  const unsigned int             dof_no_fine,
+  const MatrixFree<dim, Number> &matrix_free_coarse,
+  const unsigned int             dof_no_coarse)
+{
+  host_transfer.reinit(matrix_free_fine,
+                       dof_no_fine,
+                       matrix_free_coarse,
+                       dof_no_coarse);
+}
+
+
+template <int dim, typename VectorType>
+bool
+MGTwoLevelTransferCopyToHost<dim, VectorType>::
+  fast_polynomial_transfer_supported(const unsigned int fe_degree_fine,
+                                     const unsigned int fe_degree_coarse)
+{
+  return host_transfer.fast_polynomial_transfer_supported(fe_degree_fine,
+                                                          fe_degree_coarse);
+}
+
+template <int dim, typename VectorType>
+void
+MGTwoLevelTransferCopyToHost<dim, VectorType>::prolongate_and_add(
+  VectorType       &dst,
+  const VectorType &src) const
+{
+  if (host_vector_coarse.size() == 0)
+    host_vector_coarse.reinit(src.get_partitioner());
+  if (host_vector_fine.size() == 0)
+    host_vector_fine.reinit(dst.get_partitioner());
+
+  copy_to_host(host_vector_fine, dst);
+  copy_to_host(host_vector_coarse, src);
+
+  host_transfer.prolongate_and_add(host_vector_fine, host_vector_coarse);
+
+  copy_from_host(dst, host_vector_fine);
+}
+
+template <int dim, typename VectorType>
+void
+MGTwoLevelTransferCopyToHost<dim, VectorType>::restrict_and_add(
+  VectorType       &dst,
+  const VectorType &src) const
+{
+  if (host_vector_coarse.size() == 0)
+    host_vector_coarse.reinit(dst.get_partitioner());
+  if (host_vector_fine.size() == 0)
+    host_vector_fine.reinit(src.get_partitioner());
+
+  copy_to_host(host_vector_coarse, dst);
+  copy_to_host(host_vector_fine, src);
+
+  host_transfer.restrict_and_add(host_vector_coarse, host_vector_fine);
+
+  copy_from_host(dst, host_vector_coarse);
+}
+
+template <int dim, typename VectorType>
+void
+MGTwoLevelTransferCopyToHost<dim, VectorType>::interpolate(
+  VectorType       &dst,
+  const VectorType &src) const
+{
+  if (host_vector_coarse.size() == 0)
+    host_vector_coarse.reinit(dst.get_partitioner());
+  if (host_vector_fine.size() == 0)
+    host_vector_fine.reinit(src.get_partitioner());
+
+  copy_to_host(host_vector_fine, src);
+
+  host_transfer.interpolate(host_vector_coarse, host_vector_fine);
+
+  copy_from_host(dst, host_vector_coarse);
+}
+
+template <int dim, typename VectorType>
+std::pair<bool, bool>
+MGTwoLevelTransferCopyToHost<dim, VectorType>::
+  enable_inplace_operations_if_possible(
+    const std::shared_ptr<const Utilities::MPI::Partitioner>
+      &partitioner_coarse,
+    const std::shared_ptr<const Utilities::MPI::Partitioner> &partitioner_fine)
+{
+  return host_transfer.enable_inplace_operations_if_possible(partitioner_coarse,
+                                                             partitioner_fine);
+}
+
+template <int dim, typename VectorType>
+std::size_t
+MGTwoLevelTransferCopyToHost<dim, VectorType>::memory_consumption() const
+{
+  return host_transfer.memory_consumption();
+}
+
+template <int dim, typename VectorType>
+void
+MGTwoLevelTransferCopyToHost<dim, VectorType>::copy_to_host(
+  VectorTypeHost   &dst,
+  const VectorType &src) const
+{
+  LinearAlgebra::ReadWriteVector<Number> rw_vector(
+    src.get_partitioner()->locally_owned_range());
+  rw_vector.import_elements(src, VectorOperation::insert);
+
+  dst.import_elements(rw_vector, VectorOperation::insert);
+}
+
+template <int dim, typename VectorType>
+void
+MGTwoLevelTransferCopyToHost<dim, VectorType>::copy_from_host(
+  VectorType           &dst,
+  const VectorTypeHost &src) const
+{
+  LinearAlgebra::ReadWriteVector<Number> rw_vector(
+    src.get_partitioner()->locally_owned_range());
+  rw_vector.import_elements(src, VectorOperation::insert);
+
+  dst.import_elements(rw_vector, VectorOperation::insert);
+}
+
 
 
 template <int dim, typename Number, typename MemorySpace>