Merge pull request #1102 from brian-kelley/FixSpaddColumns

Fix C's numCols in spadd simplified interface
kokkos · Sep 13, 2021 · cf32a38 · cf32a38
2 parents 17574ed + 9f5887c
commit cf32a38
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 5 deletions.
diff --git a/src/sparse/KokkosSparse_spadd.hpp b/src/sparse/KokkosSparse_spadd.hpp
@@ -912,7 +912,6 @@ template <typename KernelHandle, typename AMatrix, typename BMatrix,
           typename CMatrix>
 void spadd_symbolic(KernelHandle* handle, const AMatrix& A, const BMatrix& B,
                     CMatrix& C) {
-  using graph_type   = typename CMatrix::staticcrsgraph_type;
   using row_map_type = typename CMatrix::row_map_type::non_const_type;
   using entries_type = typename CMatrix::index_type::non_const_type;
   using values_type  = typename CMatrix::values_type::non_const_type;
@@ -934,15 +933,12 @@ void spadd_symbolic(KernelHandle* handle, const AMatrix& A, const BMatrix& B,
   auto addHandle = handle->get_spadd_handle();
   entries_type entriesC(Kokkos::view_alloc(Kokkos::WithoutInitializing, "entries"),
                         addHandle->get_c_nnz());
-  graph_type graphC(entriesC, row_mapC);
-  C = CMatrix("matrix", graphC);
-
   // Finally since we already have the number of nnz handy
   // we can go ahead and allocate C's values and set them.
   values_type valuesC(Kokkos::view_alloc(Kokkos::WithoutInitializing, "values"),
                       addHandle->get_c_nnz());
 
-  C.values = valuesC;
+  C = CMatrix("matrix", A.numRows(), A.numCols(), addHandle->get_c_nnz(), valuesC, row_mapC, entriesC);
 }
 
 // Symbolic: count entries in each row in C to produce rowmap

diff --git a/unit_test/sparse/Test_Sparse_spadd.hpp b/unit_test/sparse/Test_Sparse_spadd.hpp
@@ -192,12 +192,58 @@ void test_spadd(lno_t numRows, lno_t numCols, size_type minNNZ, size_type maxNNZ
   }
 }
 
+//Test spadd simplified interface: make sure C's dimensions match A and B, even when there are empty rows/cols
+template <typename scalar_t, typename lno_t, typename size_type, class Device>
+void test_spadd_known_columns()
+{
+  using crsMat_t = typename KokkosSparse::CrsMatrix<scalar_t, lno_t, Device, void, size_type>;
+  using row_map_type = typename crsMat_t::row_map_type::non_const_type;
+  using entries_type = typename crsMat_t::index_type::non_const_type;
+  using values_type = typename crsMat_t::values_type::non_const_type;
+  using KAT = Kokkos::ArithTraits<scalar_t>;
+  using KernelHandle = typename KokkosKernels::Experimental::KokkosKernelsHandle<size_type, lno_t, scalar_t,
+        typename Device::execution_space, typename Device::memory_space, typename Device::memory_space>;
+  //Create A and B as 4x4 identity matrix, at the top-left of a 6x7 matrix of zeros
+  int nrows = 6;
+  int ncols = 7;
+  row_map_type Arowmap("rowmap", nrows + 1);
+  entries_type Aentries("rowmap", 4);
+  values_type Avalues("rowmap", 4);
+  {
+    auto rowmapHost = Kokkos::create_mirror_view(Arowmap);
+    auto entriesHost = Kokkos::create_mirror_view(Aentries);
+    auto valuesHost = Kokkos::create_mirror_view(Avalues);
+    for(int i = 0; i < 5; i++)
+      rowmapHost(i) = i;
+    for(int i = 5; i < nrows + 1; i++)
+      rowmapHost(i) = rowmapHost(i - 1);
+    for(int i = 0; i < 4; i++)
+    {
+      entriesHost(i) = i;
+      valuesHost(i) = KAT::one();
+    }
+    Kokkos::deep_copy(Arowmap, rowmapHost);
+    Kokkos::deep_copy(Aentries, entriesHost);
+    Kokkos::deep_copy(Avalues, valuesHost);
+  }
+  crsMat_t A("A", nrows, ncols, 4, Avalues, Arowmap, Aentries);
+  crsMat_t C;
+  KernelHandle kh;
+  kh.create_spadd_handle(true);
+  KokkosSparse::spadd_symbolic(&kh, A, A, C);
+  KokkosSparse::spadd_numeric(&kh, KAT::one(), A, KAT::one(), A, C);
+  ASSERT_EQ(nrows, C.numRows());
+  ASSERT_EQ(ncols, C.numCols());
+  ASSERT_EQ(A.nnz(), C.nnz());
+}
+
 #define EXECUTE_TEST(SCALAR, ORDINAL, OFFSET, DEVICE) \
 TEST_F( TestCategory,sparse ## _ ## spadd_sorted_input ## _ ## SCALAR ## _ ## ORDINAL ## _ ## OFFSET ## _ ## DEVICE ) { \
   test_spadd<SCALAR,ORDINAL,OFFSET,DEVICE> (10, 10, 0, 0, true); \
   test_spadd<SCALAR,ORDINAL,OFFSET,DEVICE> (10, 10, 0, 2, true); \
   test_spadd<SCALAR,ORDINAL,OFFSET,DEVICE> (100, 100, 50, 100, true); \
   test_spadd<SCALAR,ORDINAL,OFFSET,DEVICE> (50, 50, 75, 100, true); \
+  test_spadd_known_columns<SCALAR,ORDINAL,OFFSET,DEVICE> (); \
 } \
 TEST_F( TestCategory,sparse ## _ ## spadd_unsorted_input ## _ ## SCALAR ## _ ## ORDINAL ## _ ## OFFSET ## _ ## DEVICE ) { \
   test_spadd<SCALAR,ORDINAL,OFFSET,DEVICE> (10, 10, 0, 0, false); \