From: Wolfgang Bangerth <bangerth@colostate.edu>
Date: Mon, 12 Feb 2024 20:56:41 +0000 (-0700)
Subject: Apply the same patch also to the set() function.
X-Git-Tag: relicensing~40^2~1
X-Git-Url: https://gitweb.dealii.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=da4eea746a178b11a6619860c62224a7a124b08f;p=dealii.git

Apply the same patch also to the set() function.
---

diff --git a/include/deal.II/lac/trilinos_tpetra_vector.h b/include/deal.II/lac/trilinos_tpetra_vector.h
index 67bb9c42ee..83ca827f5a 100644
--- a/include/deal.II/lac/trilinos_tpetra_vector.h
+++ b/include/deal.II/lac/trilinos_tpetra_vector.h
@@ -1138,48 +1138,122 @@ namespace LinearAlgebra
       // writing to this vector at all.
       Assert(!has_ghost_elements(), ExcGhostsPresent());
 
+      // First create an alias for the type of a view into our vectors.
+      // The actual type is declared through several re-directions in
+      // Tpetra, so instead of spelling it out, we get it via decltype:
 #  if DEAL_II_TRILINOS_VERSION_GTE(13, 2, 0)
-      auto vector_2d = vector->template getLocalView<Kokkos::HostSpace>(
-        Tpetra::Access::ReadWrite);
+      using ViewType2d =
+        decltype(vector->template getLocalView<Kokkos::HostSpace>(
+          Tpetra::Access::ReadWrite));
+      ViewType2d vector_2d_local =
+        vector->template getLocalView<Kokkos::HostSpace>(
+          Tpetra::Access::ReadWrite);
 #  else
       vector->template sync<Kokkos::HostSpace>();
-      auto vector_2d = vector->template getLocalView<Kokkos::HostSpace>();
+
+      using ViewType2d =
+        decltype(vector->template getLocalView<Kokkos::HostSpace>());
+      ViewType2d vector_2d_local =
+        vector->template getLocalView<Kokkos::HostSpace>();
 #  endif
-      auto vector_1d = Kokkos::subview(vector_2d, Kokkos::ALL(), 0);
+
+      // Having extracted a view into the multivectors above, now also
+      // extract a view into the one vector we actually store. We can
+      // do this right away for the locally owned part. We defer creating
+      // the view into the nonlocal part to when we know that we actually
+      // need it; this also makes sure that we correctly deal with the
+      // case where we do not actually store a nonlocal part.
+      using ViewType1d =
+        decltype(Kokkos::subview(vector_2d_local, Kokkos::ALL(), 0));
+
+      ViewType1d vector_1d_local =
+        Kokkos::subview(vector_2d_local, Kokkos::ALL(), 0);
+      std::optional<ViewType1d> vector_1d_nonlocal;
+
 #  if !DEAL_II_TRILINOS_VERSION_GTE(13, 2, 0)
+      // Mark vector as to-be-modified. We may do the same with
+      // the nonlocal part too if we end up writing into it.
       vector->template modify<Kokkos::HostSpace>();
 #  endif
 
       for (size_type i = 0; i < n_elements; ++i)
         {
-          const size_type                   row = indices[i];
-          TrilinosWrappers::types::int_type local_row =
-            vector->getMap()->getLocalElement(row);
+          const size_type row = indices[i];
+
+          // Check if the index is in the locally owned index set.
+          // If so, we can write right into the locally owned
+          // part of the vector.
+          if (TrilinosWrappers::types::int_type local_row =
+                vector->getMap()->getLocalElement(row);
+              local_row != Teuchos::OrdinalTraits<int>::invalid())
+            {
+              vector_1d_local(local_row) = values[i];
+            }
+          else
+            {
+              // If the element was not in the locally owned part,
+              // we need to figure out whether it is in the nonlocal
+              // part. It better be:
+              Assert(nonlocal_vector.get() != nullptr, ExcInternalError());
+              TrilinosWrappers::types::int_type nonlocal_row =
+                nonlocal_vector->getMap()->getLocalElement(row);
 
 #  if DEAL_II_TRILINOS_VERSION_GTE(14, 0, 0)
-          Assert(
-            local_row != Teuchos::OrdinalTraits<int>::invalid(),
-            ExcAccessToNonLocalElement(row,
-                                       vector->getMap()->getLocalNumElements(),
-                                       vector->getMap()->getMinLocalIndex(),
-                                       vector->getMap()->getMaxLocalIndex()));
+              Assert(nonlocal_row != Teuchos::OrdinalTraits<int>::invalid(),
+                     ExcAccessToNonLocalElement(
+                       row,
+                       vector->getMap()->getLocalNumElements(),
+                       vector->getMap()->getMinLocalIndex(),
+                       vector->getMap()->getMaxLocalIndex()));
+#  else
+              Assert(nonlocal_row != Teuchos::OrdinalTraits<int>::invalid(),
+                     ExcAccessToNonLocalElement(
+                       row,
+                       vector->getMap()->getNodeNumElements(),
+                       vector->getMap()->getMinLocalIndex(),
+                       vector->getMap()->getMaxLocalIndex()));
+
+#  endif
+
+              // Having asserted that it is, write into the nonlocal part.
+              // To do so, we first need to make sure that we have a view
+              // of the nonlocal part of the vectors, since we have
+              // deferred creating this view previously:
+              if (!vector_1d_nonlocal)
+                {
+#  if DEAL_II_TRILINOS_VERSION_GTE(13, 2, 0)
+                  ViewType2d vector_2d_nonlocal =
+                    nonlocal_vector->template getLocalView<Kokkos::HostSpace>(
+                      Tpetra::Access::ReadWrite);
 #  else
-          Assert(
-            local_row != Teuchos::OrdinalTraits<int>::invalid(),
-            ExcAccessToNonLocalElement(row,
-                                       vector->getMap()->getNodeNumElements(),
-                                       vector->getMap()->getMinLocalIndex(),
-                                       vector->getMap()->getMaxLocalIndex()));
+                  ViewType2d vector_2d_nonlocal =
+                    nonlocal_vector->template getLocalView<Kokkos::HostSpace>();
 #  endif
 
-          if (local_row != Teuchos::OrdinalTraits<int>::invalid())
-            vector_1d(local_row) = values[i];
+                  vector_1d_nonlocal =
+                    Kokkos::subview(vector_2d_nonlocal, Kokkos::ALL(), 0);
+
+#  if !DEAL_II_TRILINOS_VERSION_GTE(13, 2, 0)
+                  // Mark the nonlocal vector as to-be-modified as well.
+                  nonlocal_vector->template modify<Kokkos::HostSpace>();
+#  endif
+                }
+              (*vector_1d_nonlocal)(nonlocal_row) = values[i];
+              compressed                          = false;
+            }
         }
 
 #  if !DEAL_II_TRILINOS_VERSION_GTE(13, 2, 0)
       vector->template sync<
         typename Tpetra::Vector<Number, int, types::signed_global_dof_index>::
           device_type::memory_space>();
+
+      // If we have created a view to the nonlocal part, then we have also
+      // written into it. Flush these modifications.
+      if (vector_1d_nonlocal)
+        nonlocal_vector->template sync<
+          typename Tpetra::Vector<Number, int, types::signed_global_dof_index>::
+            device_type::memory_space>();
 #  endif
     }