Optimization of export_to_ghosted_array_start()

When use CUDA-aware MPI, the sets of local indices that are ghost on other
processors are expanded in arrays. This allows for a sharp reduction in the
number of kernel launches in export_to_ghosted_array_start().

diff --git a/include/deal.II/base/partitioner.h b/include/deal.II/base/partitioner.h
index 440ada19413034da8a6d681a06cb3441402c7f9a..2c0db0955f331a2c96b8bba146a10b785cfdee4f 100644 (file)
--- a/include/deal.II/base/partitioner.h
+++ b/include/deal.II/base/partitioner.h
@@ -614,6 +614,20 @@ namespace Utilities
        */
       std::vector<std::pair<unsigned int, unsigned int>> import_indices_data;
 
+      /**
+       * The set of (local) indices that we are importing during compress(),
+       * i.e., others' ghosts that belong to the local range. The data stored is
+       * the same than in import_indices_data but the data is expanded in plain
+       * arrays. This variable is only used when using CUDA-aware MPI.
+       */
+      // The variable is mutable to enable lazy initialization in
+      // export_to_ghosted_array_start(). This way partitioner does not have to
+      // be templated on the MemorySpaceType.
+      mutable std::vector<
+        std::pair<std::unique_ptr<unsigned int[], void (*)(unsigned int *)>,
+                  unsigned int>>
+        import_indices_plain_dev;
+
       /**
        * A variable caching the number of ghost indices. It would be expensive
        * to compute it by iterating over the import indices and accumulate them.

diff --git a/include/deal.II/base/partitioner.templates.h b/include/deal.II/base/partitioner.templates.h
index f21e4e65c2a622c5572ea10329bab4b8c27e9803..ddf531bd221164fdbc7077bf5ef4a3a508ded0ff 100644 (file)
--- a/include/deal.II/base/partitioner.templates.h
+++ b/include/deal.II/base/partitioner.templates.h
@@ -98,41 +98,97 @@ namespace Utilities
         }
 
       Number *temp_array_ptr = temporary_storage.data();
-      for (unsigned int i = 0; i < n_import_targets; i++)
-        {
-          // copy the data to be sent to the import_data field
-          std::vector<std::pair<unsigned int, unsigned int>>::const_iterator
-            my_imports = import_indices_data.begin() +
-                         import_indices_chunks_by_rank_data[i],
-            end_my_imports = import_indices_data.begin() +
-                             import_indices_chunks_by_rank_data[i + 1];
-          unsigned int index = 0;
-          for (; my_imports != end_my_imports; ++my_imports)
-            {
-              const unsigned int chunk_size =
-                my_imports->second - my_imports->first;
 #    if defined(DEAL_II_COMPILER_CUDA_AWARE) && \
       defined(DEAL_II_WITH_CUDA_AWARE_MPI)
-              if (std::is_same<MemorySpaceType, MemorySpace::CUDA>::value)
+      // When using CUDAs-aware MPI, the set of local indices that are ghosts
+      // indices on other processors is expanded in arrays. This is for
+      // performance reasons as this can significantly decrease the number of
+      // kernel launched. The indices are expanded the first time the function
+      // is called.
+      if (std::is_same<MemorySpaceType, MemorySpace::CUDA>::value)
+        {
+          if (import_indices_plain_dev.size() == 0)
+            {
+              import_indices_plain_dev.reserve(n_import_targets);
+              for (unsigned int i = 0; i < n_import_targets; i++)
                 {
-                  const cudaError_t cuda_error_code =
-                    cudaMemcpy(temp_array_ptr + index,
-                               locally_owned_array.data() + my_imports->first,
-                               chunk_size * sizeof(Number),
-                               cudaMemcpyDeviceToDevice);
-                  AssertCuda(cuda_error_code);
+                  // Expand the indices on the host
+                  std::vector<std::pair<unsigned int, unsigned int>>::
+                    const_iterator my_imports =
+                                     import_indices_data.begin() +
+                                     import_indices_chunks_by_rank_data[i],
+                                   end_my_imports =
+                                     import_indices_data.begin() +
+                                     import_indices_chunks_by_rank_data[i + 1];
+                  std::vector<unsigned int> import_indices_plain_host;
+                  for (; my_imports != end_my_imports; ++my_imports)
+                    {
+                      const unsigned int chunk_size =
+                        my_imports->second - my_imports->first;
+                      for (unsigned int j = 0; j < chunk_size; ++j)
+                        import_indices_plain_host.push_back(my_imports->first +
+                                                            j);
+                    }
+
+                  // Move the indices to the device
+                  import_indices_plain_dev.emplace_back(std::make_pair(
+                    std::unique_ptr<unsigned int[], void (*)(unsigned int *)>(
+                      nullptr,
+                      Utilities::CUDA::delete_device_data<unsigned int>),
+                    import_indices_plain_host.size()));
+
+                  import_indices_plain_dev[i].first.reset(
+                    Utilities::CUDA::allocate_device_data<unsigned int>(
+                      import_indices_plain_dev[i].second));
+                  Utilities::CUDA::copy_to_dev(
+                    import_indices_plain_host,
+                    import_indices_plain_dev[i].first.get());
                 }
-              else
+            }
+        }
+#    endif
+
+      for (unsigned int i = 0; i < n_import_targets; i++)
+        {
+#    if defined(DEAL_II_COMPILER_CUDA_AWARE) && \
+      defined(DEAL_II_WITH_CUDA_AWARE_MPI)
+          if (std::is_same<MemorySpaceType, MemorySpace::CUDA>::value)
+            {
+              const int n_blocks =
+                1 + (import_indices_plain_dev[i].second - 1) /
+                      (::dealii::CUDAWrappers::chunk_size *
+                       ::dealii::CUDAWrappers::block_size);
+              ::dealii::LinearAlgebra::CUDAWrappers::kernel::
+                gather<<<n_blocks, ::dealii::CUDAWrappers::block_size>>>(
+                  temp_array_ptr,
+                  locally_owned_array.data(),
+                  import_indices_plain_dev[i].first.get(),
+                  import_indices_plain_dev[i].second);
+            }
+          else
 #    endif
+            {
+              // copy the data to be sent to the import_data field
+              std::vector<std::pair<unsigned int, unsigned int>>::const_iterator
+                my_imports = import_indices_data.begin() +
+                             import_indices_chunks_by_rank_data[i],
+                end_my_imports = import_indices_data.begin() +
+                                 import_indices_chunks_by_rank_data[i + 1];
+              unsigned int index = 0;
+              for (; my_imports != end_my_imports; ++my_imports)
                 {
-                  std::memcpy(temp_array_ptr + index,
-                              locally_owned_array.data() + my_imports->first,
-                              chunk_size * sizeof(Number));
+                  const unsigned int chunk_size =
+                    my_imports->second - my_imports->first;
+                  {
+                    std::memcpy(temp_array_ptr + index,
+                                locally_owned_array.data() + my_imports->first,
+                                chunk_size * sizeof(Number));
+                  }
+                  index += chunk_size;
                 }
-              index += chunk_size;
-            }
 
-          AssertDimension(index, import_targets_data[i].second);
+              AssertDimension(index, import_targets_data[i].second);
+            }
 
           // start the send operations
           const int ierr =
@@ -702,7 +758,7 @@ namespace Utilities
 
   } // end of namespace MPI
 
-} // end of namespace Utilities
+} // namespace Utilities
 
 
 DEAL_II_NAMESPACE_CLOSE