From 15315a7d09e6b9f304c2c95bc8fa89ccf21709df Mon Sep 17 00:00:00 2001
From: Brian Kelley <bmkelle@sandia.gov>
Date: Thu, 6 May 2021 16:24:43 -0600
Subject: [PATCH] Make CRS sorting utils work with unmanaged

(Fixing #960)
---
 src/common/KokkosKernels_Sorting.hpp     | 17 ++++---
 unit_test/common/Test_Common_Sorting.hpp | 57 ++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/src/common/KokkosKernels_Sorting.hpp b/src/common/KokkosKernels_Sorting.hpp
index 49b5dc3430..fb0d1aeeb9 100644
--- a/src/common/KokkosKernels_Sorting.hpp
+++ b/src/common/KokkosKernels_Sorting.hpp
@@ -140,15 +140,18 @@ struct SortCrsMatrixFunctor
   using lno_t = typename entries_t::non_const_value_type;
   using scalar_t = typename values_t::non_const_value_type;
   using team_mem = typename Kokkos::TeamPolicy<execution_space>::member_type;
+  //The functor owns memory for entriesAux, so it can't have MemoryTraits<Unmanaged>
+  using entries_managed_t = Kokkos::View<typename entries_t::data_type, typename entries_t::device_type>;
+  using values_managed_t = Kokkos::View<typename values_t::data_type, typename values_t::device_type>;
 
   SortCrsMatrixFunctor(bool usingRangePol, const rowmap_t& rowmap_, const entries_t& entries_, const values_t& values_)
     : rowmap(rowmap_), entries(entries_), values(values_)
   {
     if(usingRangePol)
     {
-      entriesAux = entries_t(Kokkos::ViewAllocateWithoutInitializing("Entries aux"),
+      entriesAux = entries_managed_t(Kokkos::ViewAllocateWithoutInitializing("Entries aux"),
           entries.extent(0));
-      valuesAux = values_t(Kokkos::ViewAllocateWithoutInitializing("Values aux"),
+      valuesAux = values_managed_t(Kokkos::ViewAllocateWithoutInitializing("Values aux"),
           values.extent(0));
     }
     //otherwise, aux arrays won't be allocated (sorting in place)
@@ -180,9 +183,9 @@ struct SortCrsMatrixFunctor
 
   rowmap_t rowmap;
   entries_t entries;
-  entries_t entriesAux;
+  entries_managed_t entriesAux;
   values_t values;
-  values_t valuesAux;
+  values_managed_t valuesAux;
 };
 
 template<typename execution_space, typename rowmap_t, typename entries_t>
@@ -191,13 +194,15 @@ struct SortCrsGraphFunctor
   using size_type = typename rowmap_t::non_const_value_type;
   using lno_t = typename entries_t::non_const_value_type;
   using team_mem = typename Kokkos::TeamPolicy<execution_space>::member_type;
+  //The functor owns memory for entriesAux, so it can't have MemoryTraits<Unmanaged>
+  using entries_managed_t = Kokkos::View<typename entries_t::data_type, typename entries_t::device_type>;
 
   SortCrsGraphFunctor(bool usingRangePol, const rowmap_t& rowmap_, const entries_t& entries_)
     : rowmap(rowmap_), entries(entries_)
   {
     if(usingRangePol)
     {
-      entriesAux = entries_t(Kokkos::ViewAllocateWithoutInitializing("Entries aux"),
+      entriesAux = entries_managed_t(Kokkos::ViewAllocateWithoutInitializing("Entries aux"),
           entries.extent(0));
     }
     //otherwise, aux arrays won't be allocated (sorting in place)
@@ -228,7 +233,7 @@ struct SortCrsGraphFunctor
 
   rowmap_t rowmap;
   entries_t entries;
-  entries_t entriesAux;
+  entries_managed_t entriesAux;
 };
 
 template<typename rowmap_t, typename entries_t>
diff --git a/unit_test/common/Test_Common_Sorting.hpp b/unit_test/common/Test_Common_Sorting.hpp
index b6e9de00b4..f936c7f454 100644
--- a/unit_test/common/Test_Common_Sorting.hpp
+++ b/unit_test/common/Test_Common_Sorting.hpp
@@ -642,6 +642,61 @@ void testSortCRS(default_lno_t numRows, default_lno_t numCols, default_size_type
   }
 }
 
+template<typename exec_space>
+void testSortCRSUnmanaged(bool doValues, bool doStructInterface)
+{
+  //This test is about bug #960.
+  using scalar_t = default_scalar;
+  using lno_t = default_lno_t;
+  using size_type = default_size_type;
+  using mem_space = typename exec_space::memory_space;
+  using device_t = Kokkos::Device<exec_space, mem_space>;
+  using crsMat_t = KokkosSparse::CrsMatrix<scalar_t, lno_t, device_t,
+        Kokkos::MemoryTraits<Kokkos::Unmanaged>, size_type>;
+  using crsMat_Managed_t = KokkosSparse::CrsMatrix<scalar_t, lno_t, device_t, void, size_type>;
+  using rowmap_t = typename crsMat_t::row_map_type;
+  using entries_t = typename crsMat_t::index_type;
+  using values_t = typename crsMat_t::values_type;
+  const lno_t numRows = 50;
+  const lno_t numCols = numRows;
+  size_type nnz = numRows * 5;
+  //Create a random matrix on device
+  //IMPORTANT: kk_generate_sparse_matrix does not sort the rows, if it did this
+  //wouldn't test anything
+  crsMat_Managed_t A_managed = KokkosKernels::Impl::kk_generate_sparse_matrix<crsMat_Managed_t>
+    (numRows, numCols, nnz, 2, numCols / 2);
+  crsMat_t A(A_managed);
+  auto rowmap = A.graph.row_map;
+  auto entries = A.graph.entries;
+  auto values = A.values;
+  if(doValues)
+  {
+    if(doStructInterface)
+    {
+      KokkosKernels::sort_crs_matrix(A);
+    }
+    else
+    {
+      KokkosKernels::sort_crs_matrix
+        <exec_space, rowmap_t, entries_t, values_t>
+        (A.graph.row_map, A.graph.entries, A.values);
+    }
+  }
+  else
+  {
+    if(doStructInterface)
+    {
+      KokkosKernels::sort_crs_graph(A.graph);
+    }
+    else
+    {
+      KokkosKernels::sort_crs_graph
+        <exec_space, rowmap_t, entries_t>
+        (A.graph.row_map, A.graph.entries);
+    }
+  }
+}
+
 template<typename exec_space>
 void testSortAndMerge()
 {
@@ -793,6 +848,7 @@ TEST_F( TestCategory, common_sort_crsgraph) {
     testSortCRS<TestExecSpace>(10, 10, 20, false, doStructInterface);
     testSortCRS<TestExecSpace>(100, 100, 2000, false, doStructInterface);
     testSortCRS<TestExecSpace>(1000, 1000, 30000, false, doStructInterface);
+    testSortCRSUnmanaged<TestExecSpace>(false, doStructInterface);
   }
 }
 
@@ -802,6 +858,7 @@ TEST_F( TestCategory, common_sort_crsmatrix) {
     testSortCRS<TestExecSpace>(10, 10, 20, true, doStructInterface);
     testSortCRS<TestExecSpace>(100, 100, 2000, true, doStructInterface);
     testSortCRS<TestExecSpace>(1000, 1000, 30000, true, doStructInterface);
+    testSortCRSUnmanaged<TestExecSpace>(true, doStructInterface);
   }
 }