Improvements based on copilot review

Signed-off-by: Brian Kelley <bmkelle@sandia.gov>

Author: brian-kelley

docs/source/API/sparse/spgemm_numeric.rst

@@ -50,7 +50,7 @@ Parameters
    For all parameters that are shared with ``spgemm_symbolic`` (dimensions, row maps, entries and matrices), the parameters to ``spgemm_numeric`` must match what was previously passed into ``spgemm_symbolic``
    with the same handle. For example, it is not valid to call ``spgemm_symbolic`` with one version of ``entriesA``, and then call ``spgemm_numeric`` with a different ``entriesA`` but the same handle.
 
-:handle: KokkosKernelsHandle with an active spgemm subhandle (i.e. `create_spgem_handle(...)` was previously called on it). Must be the same handle that was previously passed into ``spgemm_symbolic`` with the same input matrices.
+:handle: KokkosKernelsHandle with an active spgemm subhandle (i.e. `create_spgemm_handle(...)` was previously called on it). Must be the same handle that was previously passed into ``spgemm_symbolic`` with the same input matrices.
 
 :m, n, k: dimensions of the matrices. **m** is the number of rows in ``A`` and ``C``. **n** is the number of columns in ``A`` and the number of rows in ``B``. **k** is the number of columns in ``B`` and ``C``.
 

docs/source/API/sparse/spgemm_symbolic.rst

@@ -44,7 +44,7 @@ After ``spgemm_symbolic`` returns, the number of non-zeros in C is available thr
 Parameters
 ==========
 
-:handle: KokkosKernelsHandle with an active spgemm subhandle (i.e. `create_spgem_handle(...)` was previously called on it).
+:handle: KokkosKernelsHandle with an active spgemm subhandle (i.e. `create_spgemm_handle(...)` was previously called on it).
 
 :m, n, k: dimensions of the matrices. **m** is the number of rows in ``A`` and ``C``. **n** is the number of columns in ``A`` and the number of rows in ``B``. **k** is the number of columns in ``B`` and ``C``.
 

sparse/src/KokkosSparse_spgemm_handle.hpp

@@ -518,8 +518,7 @@ class SPGEMMHandle {
    * \param alg the SpGEMM algorithm to use.
    * \param input_sorted_ whether inputs A and B are sorted.
    *   Defaults to false for backwards compatibility.
-   * \param result_sorted_ whether the output matrix C are required to
-   *   be sorted (ascending by column index).
+   * \param result_sorted_ whether the output matrix C is required to be sorted.
    *   Defaults to true to maintain backwards compatibility.
    */
   SPGEMMHandle(SPGEMMAlgorithm alg = SPGEMM_DEFAULT, bool input_sorted_ = false, bool result_sorted_ = true)

sparse/tpls/KokkosSparse_spgemm_noreuse_tpl_spec_decl.hpp

@@ -152,7 +152,7 @@ SPGEMM_NOREUSE_DECL_CUSPARSE_S(Kokkos::complex<double>)
 
 #ifdef KOKKOSKERNELS_ENABLE_TPL_MKL
 template <typename Matrix, typename MatrixConst>
-Matrix spgemm_noreuse_mkl(const MatrixConst &A, const MatrixConst &B) {
+Matrix spgemm_noreuse_mkl(const MatrixConst &A, const MatrixConst &B, bool result_sorted) {
   using size_type   = typename Matrix::non_const_size_type;
   using index_type  = typename Matrix::non_const_ordinal_type;
   using scalar_type = typename Matrix::non_const_value_type;
@@ -171,7 +171,9 @@ Matrix spgemm_noreuse_mkl(const MatrixConst &A, const MatrixConst &B) {
   generalDescr.mode = SPARSE_FILL_MODE_FULL;
   generalDescr.diag = SPARSE_DIAG_NON_UNIT;
   KOKKOSKERNELS_MKL_SAFE_CALL(mkl_sparse_spmm(SPARSE_OPERATION_NON_TRANSPOSE, Amkl, Bmkl, &C));
-  KOKKOSKERNELS_MKL_SAFE_CALL(mkl_sparse_order(C));
+  if (result_sorted) {
+    KOKKOSKERNELS_MKL_SAFE_CALL(mkl_sparse_order(C));
+  }
   MKLMatrix wrappedC(C);
   MKL_INT nrows = 0, ncols = 0;
   MKL_INT *rowmapRaw     = nullptr;
@@ -216,10 +218,10 @@ Matrix spgemm_noreuse_mkl(const MatrixConst &A, const MatrixConst &B) {
                                                 Kokkos::MemoryTraits<Kokkos::Unmanaged>, const MKL_INT>;              \
     static KokkosSparse::CrsMatrix<SCALAR, MKL_INT, Kokkos::Device<EXEC, Kokkos::HostSpace>, void, MKL_INT>           \
     spgemm_noreuse(KokkosSparse::SPGEMMAlgorithm, const ConstMatrix &A, bool, const ConstMatrix &B, bool, bool,       \
-                   bool) {                                                                                            \
+                   bool result_sorted) {                                                                              \
       std::string label = "KokkosSparse::spgemm_noreuse[TPL_MKL," + KokkosKernels::ArithTraits<SCALAR>::name() + "]"; \
       Kokkos::Profiling::pushRegion(label);                                                                           \
-      Matrix C = spgemm_noreuse_mkl<Matrix>(A, B);                                                                    \
+      Matrix C = spgemm_noreuse_mkl<Matrix>(A, B, result_sorted);                                                     \
       Kokkos::Profiling::popRegion();                                                                                 \
       return C;                                                                                                       \
     }                                                                                                                 \

sparse/unit_test/Test_Sparse_Utils.hpp

@@ -23,6 +23,8 @@ vector_t create_random_y_vector_mv(crsMat_t crsMat, vector_t x_vector) {
   return y_vector;
 }
 
+// Test whether two CRS matrices are identical (shape, #nonzeros, entries, and values to within tolerance),
+// including the order of entries within each row.
 template <typename crsMat_t, typename device>
 bool is_same_matrix(crsMat_t output_mat_actual, crsMat_t output_mat_reference) {
   typedef typename crsMat_t::StaticCrsGraphType graph_t;

sparse/unit_test/Test_Sparse_spgemm.hpp

@@ -56,12 +56,14 @@ void randomize_matrix_values(const Values &v) {
 }
 
 template <typename crsMat_t>
-void run_spgemm_noreuse(KokkosSparse::SPGEMMAlgorithm algo, crsMat_t A, crsMat_t B, crsMat_t &C) {
-  C = KokkosSparse::spgemm<crsMat_t>(algo, A, false, B, false);
+void run_spgemm_noreuse(KokkosSparse::SPGEMMAlgorithm algo, crsMat_t A, crsMat_t B, crsMat_t &C, bool input_sorted,
+                        bool result_sorted) {
+  C = KokkosSparse::spgemm<crsMat_t>(algo, A, false, B, false, input_sorted, result_sorted);
 }
 
 template <typename crsMat_t, typename device>
-int run_spgemm(crsMat_t &A, crsMat_t &B, KokkosSparse::SPGEMMAlgorithm spgemm_algorithm, crsMat_t &C, bool testReuse) {
+int run_spgemm(crsMat_t &A, crsMat_t &B, KokkosSparse::SPGEMMAlgorithm spgemm_algorithm, crsMat_t &C, bool testReuse,
+               bool input_sorted, bool result_sorted) {
   typedef typename crsMat_t::size_type size_type;
   typedef typename crsMat_t::ordinal_type lno_t;
   typedef typename crsMat_t::value_type scalar_t;
@@ -75,7 +77,7 @@ int run_spgemm(crsMat_t &A, crsMat_t &B, KokkosSparse::SPGEMMAlgorithm spgemm_al
   kh.set_team_work_size(16);
   kh.set_dynamic_scheduling(true);
 
-  kh.create_spgemm_handle(spgemm_algorithm);
+  kh.create_spgemm_handle(spgemm_algorithm, input_sorted, result_sorted);
   {
     auto sh = kh.get_spgemm_handle();
 
@@ -112,7 +114,7 @@ int run_spgemm(crsMat_t &A, crsMat_t &B, KokkosSparse::SPGEMMAlgorithm spgemm_al
 
 template <typename crsMat_t, typename device>
 int run_spgemm_old_interface(crsMat_t &A, crsMat_t &B, KokkosSparse::SPGEMMAlgorithm spgemm_algorithm, crsMat_t &result,
-                             bool testReuse) {
+                             bool testReuse, bool input_sorted, bool result_sorted) {
   typedef typename crsMat_t::StaticCrsGraphType graph_t;
   typedef typename graph_t::row_map_type::non_const_type lno_view_t;
   typedef typename graph_t::entries_type::non_const_type lno_nnz_view_t;
@@ -130,7 +132,7 @@ int run_spgemm_old_interface(crsMat_t &A, crsMat_t &B, KokkosSparse::SPGEMMAlgor
   kh.set_team_work_size(16);
   kh.set_dynamic_scheduling(true);
 
-  kh.create_spgemm_handle(spgemm_algorithm);
+  kh.create_spgemm_handle(spgemm_algorithm, input_sorted, result_sorted);
   {
     auto sh = kh.get_spgemm_handle();
 
@@ -197,14 +199,8 @@ void test_spgemm(lno_t m, lno_t k, lno_t n, size_type nnz, lno_t bandwidth, lno_
   }
 #endif  // KOKKOSKERNELS_ENABLE_TPL_ARMPL
   using namespace Test;
-  // device::execution_space::initialize();
-  // device::execution_space::print_configuration(std::cout);
 
   typedef CrsMatrix<scalar_t, lno_t, device, void, size_type> crsMat_t;
-  // typedef typename crsMat_t::StaticCrsGraphType graph_t;
-  // typedef typename graph_t::row_map_type::non_const_type lno_view_t;
-  // typedef typename graph_t::entries_type::non_const_type   lno_nnz_view_t;
-  // typedef typename crsMat_t::values_type::non_const_type scalar_view_t;
 
   // Generate random compressed sparse row matrix. Randomly generated (non-zero)
   // values are stored in a 1-D (1 rank) array.
@@ -220,7 +216,8 @@ void test_spgemm(lno_t m, lno_t k, lno_t n, size_type nnz, lno_t bandwidth, lno_
   // If this test won't reuse symbolic, we can compute the reference matrix once
   // here
   if (!testReuse) {
-    run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, output_mat2, false);
+    run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, output_mat2, /* testReuse */ false, /* input sorted */ true,
+                                 /* result sorted */ true);
   }
 
   std::vector<SPGEMMAlgorithm> algorithms;
@@ -280,12 +277,16 @@ void test_spgemm(lno_t m, lno_t k, lno_t n, size_type nnz, lno_t bandwidth, lno_
     try {
       switch (callMode) {
         case spgemm_reuse_view:
-          res = run_spgemm_old_interface<crsMat_t, device>(A, B, spgemm_algorithm, output_mat, testReuse);
+          res = run_spgemm_old_interface<crsMat_t, device>(A, B, spgemm_algorithm, output_mat, testReuse,
+                                                           /* input sorted */ true, /* result sorted */ true);
           break;
         case spgemm_reuse_matrix:
-          res = run_spgemm<crsMat_t, device>(A, B, spgemm_algorithm, output_mat, testReuse);
+          res = run_spgemm<crsMat_t, device>(A, B, spgemm_algorithm, output_mat, testReuse, /* input sorted */ true,
+                                             /* result sorted */ true);
+          break;
+        case spgemm_noreuse:
+          run_spgemm_noreuse(spgemm_algorithm, A, B, output_mat, /* input sorted */ true, /* result sorted */ true);
           break;
-        case spgemm_noreuse: run_spgemm_noreuse(spgemm_algorithm, A, B, output_mat); break;
       }
     } catch (const char *message) {
       EXPECT_TRUE(is_expected_to_fail) << algo << ": " << message;
@@ -302,7 +303,8 @@ void test_spgemm(lno_t m, lno_t k, lno_t n, size_type nnz, lno_t bandwidth, lno_
     // If this is testing reuse, the values of A and B changed so
     // the reference matrix must be recomputed
     if (testReuse) {
-      run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, output_mat2, false);
+      run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, output_mat2, /* reuse */ false, /* input sorted */ true,
+                                   /* result sorted */ true);
     }
 
     // double spgemm_time = timer1.seconds();
@@ -315,12 +317,8 @@ void test_spgemm(lno_t m, lno_t k, lno_t n, size_type nnz, lno_t bandwidth, lno_
                                 << ", testReuse=" << int(testReuse) << ", m=" << m << ", k=" << k << ", n=" << n
                                 << ", nnz=" << nnz << ", bandwidth=" << bandwidth
                                 << ", row_size_variance=" << row_size_variance;
-      // EXPECT_TRUE( equal) << algo;
     }
-    // std::cout << "algo:" << algo << " spgemm_time:" << spgemm_time << "
-    // output_check_time:" << timer1.seconds() << std::endl;
   }
-  // device::execution_space::finalize();
 }
 
 template <typename scalar_t, typename lno_t, typename size_type, typename device>
@@ -358,14 +356,17 @@ void test_spgemm_symbolic(bool callSymbolicFirst, bool testEmpty) {
     KokkosSparse::sort_crs_matrix(B);
   }
   // Call reference impl to get complete product
+  // Note: for both reference C and output C: we're only checking that
+  // the rowmap is correct, so don't care if result is sorted.
   crsMat_t C_reference;
-  Test::run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, C_reference, false);
+  Test::run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, C_reference, /* reuse */ false, /* input sorted */ true,
+                                     /* result sorted */ false);
   // Now call just symbolic, and specifically request that rowptrs be populated
   // Make sure this never depends on C_rowmap being initialized
   rowmap_t C_rowmap(Kokkos::view_alloc(Kokkos::WithoutInitializing, "rowmapC"), m + 1);
   Kokkos::deep_copy(C_rowmap, size_type(123));
   KernelHandle kh;
-  kh.create_spgemm_handle();
+  kh.create_spgemm_handle(KokkosSparse::SPGEMM_DEFAULT, /* input sorted */ true, /* output sorted */ false);
   if (callSymbolicFirst) {
     KokkosSparse::spgemm_symbolic(&kh, m, n, k, A.graph.row_map, A.graph.entries, false, B.graph.row_map,
                                   B.graph.entries, false, C_rowmap);
@@ -430,12 +431,14 @@ void test_issue402() {
   KokkosSparse::sort_crs_matrix(A);
   KokkosSparse::sort_crs_matrix(B);
   crsMat_t Cgold;
-  run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, Cgold, false);
+  run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, Cgold, /* reuse */ false, /* input sorted */ true,
+                               /* result sorted */ true);
   crsMat_t C;
   bool success = true;
   std::string errMsg;
   try {
-    int res = run_spgemm<crsMat_t, device>(A, B, SPGEMM_KK_MEMORY, C, false);
+    int res = run_spgemm<crsMat_t, device>(A, B, SPGEMM_KK_MEMORY, C, /* reuse */ false, /* input sorted */ true,
+                                           /* result sorted */ true);
     if (res) throw "run_spgemm returned error code";
   } catch (const char *message) {
     errMsg  = message;
@@ -520,18 +523,18 @@ void test_spgemm_sortedness() {
   const lno_t bandwidth         = 500;
   const lno_t row_size_variance = 10;
 
-  // Generate random matrices. This relies on kk_generate_sparse_matrix not producing a sorted matrix.
+  // Generate random matrices. This relies on kk_generate_sparse_matrix NOT producing a sorted matrix.
   crsMat_t A = KokkosSparse::Impl::kk_generate_sparse_matrix<crsMat_t>(m, k, nnz, row_size_variance, bandwidth);
   crsMat_t B = KokkosSparse::Impl::kk_generate_sparse_matrix<crsMat_t>(k, n, nnz, row_size_variance, bandwidth);
   randomize_matrix_values(A.values);
   randomize_matrix_values(B.values);
 
   // Compute reference C using the serial/debug algorithm.
   crsMat_t C_reference;
-  run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, C_reference, false);
+  run_spgemm<crsMat_t, device>(A, B, SPGEMM_DEBUG, C_reference, /* reuse */ false, /* input sorted */ false,
+                               /* result sorted */ true);
 
   std::vector<std::string> algorithm_names = {"SPGEMM_DEFAULT", "SPGEMM_KK"};
-  std::vector<SPGEMMAlgorithm> algorithms  = {SPGEMM_DEFAULT, SPGEMM_KK};
   for (auto algo_name : algorithm_names) {
     SPGEMMAlgorithm algo = KokkosSparse::StringToSPGEMMAlgorithm(algo_name);
     for (bool resultSorted : {true, false}) {
@@ -560,7 +563,7 @@ void test_spgemm_sortedness() {
           kh.destroy_spgemm_handle();
           C = crsMat_t("C", num_rows_A, num_cols_B, c_nnz_size, valuesC, row_mapC, entriesC);
         } else {
-          run_spgemm_noreuse(algo, A, B, C);
+          run_spgemm_noreuse(algo, A, B, C, /* input sorted */ false, /* result sorted */ resultSorted);
         }
 
         // If we requested sorted output, verify the output really is sorted.
@@ -570,7 +573,7 @@ void test_spgemm_sortedness() {
                               << " produced unsorted output for algo " << algo_name
                               << " even though sorted output was requested";
         } else {
-          // Sort C for comparison
+          // Sort C so we can directly compare it to sorted reference
           KokkosSparse::sort_crs_matrix(C);
         }