Tpetra MatrixMatrix: sort for cuSparse

trilinos · Sep 4, 2024 · 34840d9 · 34840d9
1 parent 202c11a
commit 34840d9
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/packages/tpetra/core/ext/TpetraExt_MatrixMatrix_Cuda.hpp b/packages/tpetra/core/ext/TpetraExt_MatrixMatrix_Cuda.hpp
@@ -170,7 +170,20 @@ void KernelWrappers<Scalar,LocalOrdinal,GlobalOrdinal,Tpetra::KokkosCompat::Kokk
   KokkosSparse::SPGEMMAlgorithm alg_enum = KokkosSparse::StringToSPGEMMAlgorithm(myalg);
 
   // Merge the B and Bimport matrices
-  const KCRS Bmerged = Tpetra::MMdetails::merge_matrices(Aview,Bview,Acol2Brow,Acol2Irow,Bcol2Ccol,Icol2Ccol,C.getColMap()->getLocalNumElements());
+  KCRS Bmerged = Tpetra::MMdetails::merge_matrices(Aview,Bview,Acol2Brow,Acol2Irow,Bcol2Ccol,Icol2Ccol,C.getColMap()->getLocalNumElements());
+
+
+  // if Kokkos Kernels is going to use the cuSparse TPL for SpGEMM, this matrix
+  // needs to be sorted
+  // Try to mirror the Kokkos Kernels internal SpGEMM TPL use logic
+  // inspired by https://github.com/trilinos/Trilinos/pull/11709
+#if defined(KOKKOS_ENABLE_CUDA) \
+ && defined(KOKKOSKERNELS_ENABLE_TPL_CUSPARSE) \
+ && ((CUDA_VERSION < 11000) || (CUDA_VERSION >= 11040))
+  if constexpr (std::is_same_v<typename device_t::execution_space, Kokkos::Cuda>) {
+    KokkosSparse::sort_crs_matrix(Bmerged);
+  }
+#endif
 
 #ifdef HAVE_TPETRA_MMM_TIMINGS
   MM = Teuchos::null; MM = rcp(new TimeMonitor (*TimeMonitor::getNewTimer(prefix_mmm + std::string("MMM Newmatrix CudaCore"))));