Ifpack2: AdditiveSchwarz matrix filtering in parallel

Add a fast path ("Details::AdditiveSchwarzFilter") where LocalFilter, SingletonFilter and ReorderFilter are applied all at once on device. This benefits initialize, compute and apply, since the old path only supported host matrix access with getLocalRowCopy, and each filter pass did its own copy. The old path is still in place to support non-CrsMatrix inputs to AdditiveSchwarz. Also fix trilinos#9606 by having AdditiveSchwarz::compute re-import the halo rows of the OverlappingRowMatrix (ExtMatrix_).
brian-kelley · Jul 1, 2022 · 85a779b · 85a779b
1 parent 0f5071f
commit 85a779b
Show file tree

Hide file tree

Showing 12 changed files with 1,460 additions and 184 deletions.
diff --git a/packages/ifpack2/src/CMakeLists.txt b/packages/ifpack2/src/CMakeLists.txt
@@ -248,6 +248,7 @@ IF(Ifpack2_ENABLE_EXPLICIT_INSTANTIATION)
     SparsityFilter
     ContainerFactory
     TriDiContainer
+    Details::AdditiveSchwarzFilter
     Details::Chebyshev
     Details::ChebyshevKernel
     Details::DenseSolver

diff --git a/packages/ifpack2/src/Ifpack2_AdditiveSchwarz_decl.hpp b/packages/ifpack2/src/Ifpack2_AdditiveSchwarz_decl.hpp
@@ -56,6 +56,7 @@
 #include "Ifpack2_Preconditioner.hpp"
 #include "Ifpack2_Details_CanChangeMatrix.hpp"
 #include "Ifpack2_Details_NestedPreconditioner.hpp"
+#include "Ifpack2_OverlappingRowMatrix.hpp"
 #include "Tpetra_Map.hpp"
 #include "Tpetra_MultiVector.hpp"
 #include "Tpetra_RowMatrix.hpp"
@@ -329,6 +330,12 @@ class AdditiveSchwarz :
   using row_matrix_type =
     Tpetra::RowMatrix<scalar_type, local_ordinal_type,
                       global_ordinal_type, node_type>;
+
+  //! The Tpetra::CrsMatrix specialization that is a subclass of MatrixType.
+  using crs_matrix_type =
+    Tpetra::CrsMatrix<scalar_type, local_ordinal_type,
+                      global_ordinal_type, node_type>;
+
   //@}
   // \name Constructors and destructor
   //@{
@@ -774,7 +781,7 @@ class AdditiveSchwarz :
   /// \brief The overlapping matrix.
   ///
   /// If nonnull, this is an instance of OverlappingRowMatrix<row_matrix_type>.
-  Teuchos::RCP<row_matrix_type> OverlappingMatrix_;
+  Teuchos::RCP<OverlappingRowMatrix<row_matrix_type>> OverlappingMatrix_;
 
   //! The reordered matrix.
   Teuchos::RCP<row_matrix_type> ReorderedLocalizedMatrix_;

diff --git a/packages/ifpack2/src/Ifpack2_AdditiveSchwarz_def.hpp b/packages/ifpack2/src/Ifpack2_AdditiveSchwarz_def.hpp
diff --git a/packages/ifpack2/src/Ifpack2_Details_AdditiveSchwarzFilter_decl.hpp b/packages/ifpack2/src/Ifpack2_Details_AdditiveSchwarzFilter_decl.hpp
diff --git a/packages/ifpack2/src/Ifpack2_Details_AdditiveSchwarzFilter_def.hpp b/packages/ifpack2/src/Ifpack2_Details_AdditiveSchwarzFilter_def.hpp
diff --git a/packages/ifpack2/src/Ifpack2_IlukGraph.hpp b/packages/ifpack2/src/Ifpack2_IlukGraph.hpp
@@ -281,6 +281,7 @@ void IlukGraph<GraphType, KKHandleType>::initialize()
 
   // Get Maximum Row length
   const int MaxNumIndices = OverlapGraph_->getLocalMaxNumRowEntries ();
+  std::cout << "IlukGraph: OverlapGraph says max num row entries is " << MaxNumIndices << '\n';
 
   // FIXME (mfh 23 Dec 2013) Use size_t or whatever
   // getLocalNumElements() returns, instead of ptrdiff_t.

diff --git a/packages/ifpack2/src/Ifpack2_OverlappingRowMatrix_decl.hpp b/packages/ifpack2/src/Ifpack2_OverlappingRowMatrix_decl.hpp
@@ -75,6 +75,8 @@ class OverlappingRowMatrix :
 
   using row_matrix_type = Tpetra::RowMatrix<scalar_type, local_ordinal_type,
 					    global_ordinal_type, node_type>;
+  using crs_matrix_type = Tpetra::CrsMatrix<scalar_type, local_ordinal_type,
+					    global_ordinal_type, node_type>;
 
   static_assert(std::is_same<MatrixType, row_matrix_type>::value, "Ifpack2::OverlappingRowMatrix: The template parameter MatrixType must be a Tpetra::RowMatrix specialization.  Please don't use Tpetra::CrsMatrix (a subclass of Tpetra::RowMatrix) here anymore.  The constructor can take either a RowMatrix or a CrsMatrix just fine.");
 
@@ -341,18 +343,19 @@ class OverlappingRowMatrix :
 
   void describe(Teuchos::FancyOStream &out, const Teuchos::EVerbosityLevel verbLevel) const;
 
-  virtual Teuchos::RCP<const row_matrix_type> getUnderlyingMatrix() const;
+  Teuchos::RCP<const crs_matrix_type> getUnderlyingMatrix() const;
 
-  Teuchos::RCP<const row_matrix_type> getExtMatrix() const;
+  Teuchos::RCP<const crs_matrix_type> getExtMatrix() const;
 
   Teuchos::ArrayView<const size_t> getExtHaloStarts() const;
 
+  void doExtImport();
+
 private:
   typedef Tpetra::Map<local_ordinal_type, global_ordinal_type, node_type> map_type;
   typedef Tpetra::Import<local_ordinal_type, global_ordinal_type, node_type> import_type;
   typedef Tpetra::Export<local_ordinal_type, global_ordinal_type, node_type> export_type;
   typedef Tpetra::RowGraph<local_ordinal_type, global_ordinal_type, node_type> row_graph_type;
-  typedef Tpetra::CrsMatrix<scalar_type, local_ordinal_type, global_ordinal_type, node_type> crs_matrix_type;
   typedef Tpetra::Vector<scalar_type, local_ordinal_type, global_ordinal_type, node_type> vector_type;
 
   //! The input matrix to the constructor.
@@ -369,7 +372,7 @@ class OverlappingRowMatrix :
   Teuchos::RCP<const import_type> Importer_;
 
   //! The matrix containing only the overlap rows.
-  Teuchos::RCP<const crs_matrix_type> ExtMatrix_;
+  Teuchos::RCP<crs_matrix_type> ExtMatrix_;
   Teuchos::RCP<const map_type>        ExtMap_;
   Teuchos::RCP<const import_type>     ExtImporter_;
   Teuchos::Array<size_t>              ExtHaloStarts_;

diff --git a/packages/ifpack2/src/Ifpack2_OverlappingRowMatrix_def.hpp b/packages/ifpack2/src/Ifpack2_OverlappingRowMatrix_def.hpp
@@ -175,11 +175,9 @@ OverlappingRowMatrix (const Teuchos::RCP<const row_matrix_type>& A,
   ExtImporter_ = rcp (new import_type (A_->getRowMap (), ExtMap_));
 
   {
-    RCP<crs_matrix_type> ExtMatrix_nc =
-      rcp (new crs_matrix_type (ExtMap_, ColMap_, 0));
-    ExtMatrix_nc->doImport (*A_, *ExtImporter_, Tpetra::INSERT);
-    ExtMatrix_nc->fillComplete (A_->getDomainMap (), RowMap_);
-    ExtMatrix_ = ExtMatrix_nc; // we only need the const version after here
+    ExtMatrix_ = rcp (new crs_matrix_type (ExtMap_, ColMap_, 0));
+    ExtMatrix_->doImport (*A_, *ExtImporter_, Tpetra::INSERT);
+    ExtMatrix_->fillComplete (A_->getDomainMap (), RowMap_);
   }
 
   // fix indices for overlapping matrix
@@ -819,14 +817,14 @@ void OverlappingRowMatrix<MatrixType>::describe(Teuchos::FancyOStream &out,
 }
 
 template<class MatrixType>
-Teuchos::RCP<const Tpetra::RowMatrix<typename MatrixType::scalar_type, typename MatrixType::local_ordinal_type, typename MatrixType::global_ordinal_type, typename MatrixType::node_type> >
+Teuchos::RCP<const Tpetra::CrsMatrix<typename MatrixType::scalar_type, typename MatrixType::local_ordinal_type, typename MatrixType::global_ordinal_type, typename MatrixType::node_type> >
 OverlappingRowMatrix<MatrixType>::getUnderlyingMatrix() const
 {
   return A_;
 }
 
 template<class MatrixType>
-Teuchos::RCP<const Tpetra::RowMatrix<typename MatrixType::scalar_type, typename MatrixType::local_ordinal_type, typename MatrixType::global_ordinal_type, typename MatrixType::node_type> >
+Teuchos::RCP<const Tpetra::CrsMatrix<typename MatrixType::scalar_type, typename MatrixType::local_ordinal_type, typename MatrixType::global_ordinal_type, typename MatrixType::node_type> >
 OverlappingRowMatrix<MatrixType>::getExtMatrix() const
 {
   return ExtMatrix_;
@@ -838,6 +836,13 @@ Teuchos::ArrayView<const size_t> OverlappingRowMatrix<MatrixType>::getExtHaloSta
   return ExtHaloStarts_();
 }
 
+template<class MatrixType>
+void OverlappingRowMatrix<MatrixType>::doExtImport()
+{
+  ExtMatrix_ = rcp (new crs_matrix_type (ExtMap_, ColMap_, 0));
+  ExtMatrix_->doImport (*A_, *ExtImporter_, Tpetra::INSERT);
+  ExtMatrix_->fillComplete (A_->getDomainMap (), RowMap_);
+}
 
 } // namespace Ifpack2
 

diff --git a/packages/ifpack2/test/unit_tests/Ifpack2_UnitTestAdditiveSchwarz.cpp b/packages/ifpack2/test/unit_tests/Ifpack2_UnitTestAdditiveSchwarz.cpp
@@ -133,9 +133,11 @@ TEUCHOS_UNIT_TEST_TEMPLATE_3_DECL(Ifpack2AdditiveSchwarz, Test0, Scalar, LocalOr
   params.set ("schwarz: combine mode", "Zero");
 
 #if defined(HAVE_IFPACK2_XPETRA) && defined(HAVE_IFPACK2_ZOLTAN2)
+  std::cout << "Test0: Enabling reordering!\n";
   params.set ("schwarz: use reordering", true);
   params.set ("schwarz: reordering list", zlist);
 #else
+  std::cout << "Test0: NOT enabling reordering!\n";
   params.set ("schwarz: use reordering", false);
 #endif
 

diff --git a/packages/ifpack2/test/unit_tests/Ifpack2_UnitTestLocalSparseTriangularSolver.cpp b/packages/ifpack2/test/unit_tests/Ifpack2_UnitTestLocalSparseTriangularSolver.cpp
@@ -84,15 +84,8 @@ localSolve (Tpetra::MultiVector<
   using Teuchos::CONJ_TRANS;
   using Teuchos::NO_TRANS;
   using Teuchos::TRANS;
-  using MV = Tpetra::MultiVector<
-    typename CrsMatrixType::scalar_type,
-    typename CrsMatrixType::local_ordinal_type,
-    typename CrsMatrixType::global_ordinal_type,
-    typename CrsMatrixType::node_type>;
   using scalar_type = typename CrsMatrixType::scalar_type;
   using STS = Teuchos::ScalarTraits<scalar_type>;
-  using device_type = typename CrsMatrixType::device_type;
-  using dev_memory_space = typename device_type::memory_space;
   const char prefix[] = "localSolve: ";
 
   TEUCHOS_TEST_FOR_EXCEPTION
@@ -999,12 +992,6 @@ void testArrowMatrix (bool& success, Teuchos::FancyOStream& out)
                                      L,U,out);
   if(!gblSuccess) return;
 
-  typedef typename crs_matrix_type::local_graph_device_type local_graph_type;
-  typedef typename crs_matrix_type::local_matrix_device_type local_matrix_type;
-  typedef typename local_matrix_type::row_map_type::non_const_type row_offsets_type;
-  typedef typename local_graph_type::entries_type::non_const_type col_inds_type;
-  typedef typename local_matrix_type::values_type::non_const_type values_type;
-
   typedef typename crs_matrix_type::local_inds_host_view_type const_local_inds_type;
   typedef typename crs_matrix_type::values_host_view_type const_values_type;
 

diff --git a/packages/tpetra/core/src/Tpetra_CrsGraph_decl.hpp b/packages/tpetra/core/src/Tpetra_CrsGraph_decl.hpp
@@ -1326,7 +1326,7 @@ namespace Tpetra {
     ///
     /// \pre columnIndices are sorted within rows
     /// \pre <tt>hasColMap() == true</tt>
-    /// \pre <tt>rowPointers.size() != getLocalNumRows()+1</tt>
+    /// \pre <tt>rowPointers.size() == getLocalNumRows()+1</tt>
     /// \pre No insert routines have been called.
     ///
     /// \warning This method is intended for expert developer use
@@ -1339,7 +1339,7 @@ namespace Tpetra {
     ///
     /// \pre columnIndices are sorted within rows
     /// \pre <tt>hasColMap() == true</tt>
-    /// \pre <tt>rowPointers.size() != getLocalNumRows()+1</tt>
+    /// \pre <tt>rowPointers.size() == getLocalNumRows()+1</tt>
     /// \pre No insert routines have been called.
     ///
     /// \warning This method is intended for expert developer use

diff --git a/packages/tpetra/core/src/Tpetra_CrsGraph_def.hpp b/packages/tpetra/core/src/Tpetra_CrsGraph_def.hpp
@@ -2720,13 +2720,15 @@ namespace Tpetra {
       }
     }
 
+    /*
     // FIXME (mfh 07 Aug 2014) We need to relax this restriction,
     // since the future model will be allocation at construction, not
     // lazy allocation on first insert.
     TEUCHOS_TEST_FOR_EXCEPTION_CLASS_FUNC
       ((this->lclIndsUnpacked_wdv.extent (0) != 0 || this->gblInds_wdv.extent (0) != 0),
        std::runtime_error, "You may not call this method if 1-D data "
        "structures are already allocated.");
+     */
 
     indicesAreAllocated_ = true;
     indicesAreLocal_     = true;