Skip to content

Commit

Permalink
Make CRS sorting utils work with unmanaged
Browse files Browse the repository at this point in the history
(Fixing #960)
  • Loading branch information
brian-kelley committed May 6, 2021
1 parent 74c71c5 commit 15315a7
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 6 deletions.
17 changes: 11 additions & 6 deletions src/common/KokkosKernels_Sorting.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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>
Expand All @@ -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)
Expand Down Expand Up @@ -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>
Expand Down
57 changes: 57 additions & 0 deletions unit_test/common/Test_Common_Sorting.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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()
{
Expand Down Expand Up @@ -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);
}
}

Expand All @@ -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);
}
}

Expand Down

0 comments on commit 15315a7

Please sign in to comment.