Sparse - sptrsv: implementing generic tpl layer for sptrsv

Adding cusparse to tpl spec avail although will need much more to actually leverage the TPL layer.
kokkos · May 8, 2024 · 2f34a69 · 2f34a69
1 parent dfacd72
commit 2f34a69
Show file tree

Hide file tree

Showing 8 changed files with 876 additions and 136 deletions.
diff --git a/sparse/impl/KokkosSparse_crs_to_bsr_impl.hpp b/sparse/impl/KokkosSparse_crs_to_bsr_impl.hpp
@@ -99,6 +99,7 @@ template <typename Bsr, typename Crs>
 Bsr blocked_crs_to_bsr(const Crs &crs, size_t blockSize) {
   using bsr_value_type   = typename Bsr::value_type;
   using bsr_ordinal_type = typename Bsr::ordinal_type;
+  using csr_size_type    = typename Crs::non_const_size_type;
 
   // copy matrix data to host
   auto hRowMap  = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),
@@ -119,7 +120,7 @@ Bsr blocked_crs_to_bsr(const Crs &crs, size_t blockSize) {
 
   for (bsr_ordinal_type row = 0; row < bsr_ordinal_type(hRowMap.size()) - 1;
        ++row) {
-    for (size_t ci = hRowMap(row); ci < hRowMap(row + 1); ++ci) {
+    for (csr_size_type ci = hRowMap(row); ci < hRowMap(row + 1); ++ci) {
       bsr_ordinal_type col = hColInds(ci);
       bsr_value_type val   = hVals(ci);
 

diff --git a/sparse/impl/KokkosSparse_sptrsv_symbolic_spec.hpp b/sparse/impl/KokkosSparse_sptrsv_symbolic_spec.hpp
@@ -68,27 +68,29 @@ namespace Impl {
 /// \brief Implementation of KokkosSparse::sptrsv_symbolic
 
 template <class ExecutionSpace, class KernelHandle, class RowMapType,
-          class EntriesType,
+          class EntriesType, class ValuesType,
           bool tpl_spec_avail = sptrsv_symbolic_tpl_spec_avail<
-              KernelHandle, RowMapType, EntriesType>::value,
+			       KernelHandle, RowMapType, EntriesType, ValuesType>::value,
           bool eti_spec_avail = sptrsv_symbolic_eti_spec_avail<
               KernelHandle, RowMapType, EntriesType>::value>
 struct SPTRSV_SYMBOLIC {
   static void sptrsv_symbolic(const ExecutionSpace &space, KernelHandle *handle,
                               const RowMapType row_map,
-                              const EntriesType entries);
+                              const EntriesType entries,
+			      const ValuesType values);
 };
 
 #if !defined(KOKKOSKERNELS_ETI_ONLY) || KOKKOSKERNELS_IMPL_COMPILE_LIBRARY
 //! Full specialization of sptrsv_symbolic
 // Unification layer
 template <class ExecutionSpace, class KernelHandle, class RowMapType,
-          class EntriesType>
-struct SPTRSV_SYMBOLIC<ExecutionSpace, KernelHandle, RowMapType, EntriesType,
+          class EntriesType, class ValuesType>
+struct SPTRSV_SYMBOLIC<ExecutionSpace, KernelHandle, RowMapType, EntriesType, ValuesType,
                        false, KOKKOSKERNELS_IMPL_COMPILE_LIBRARY> {
   static void sptrsv_symbolic(const ExecutionSpace &space, KernelHandle *handle,
                               const RowMapType row_map,
-                              const EntriesType entries) {
+                              const EntriesType entries,
+			      const ValuesType /*values*/) {
     auto sptrsv_handle = handle->get_sptrsv_handle();
     auto nrows         = row_map.extent(0) - 1;
     sptrsv_handle->new_init_handle(nrows);
@@ -129,6 +131,10 @@ struct SPTRSV_SYMBOLIC<ExecutionSpace, KernelHandle, RowMapType, EntriesType,
           const ORDINAL_TYPE *, LAYOUT_TYPE,                                \
           Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,                  \
           Kokkos::MemoryTraits<Kokkos::Unmanaged | Kokkos::RandomAccess> >, \
+      Kokkos::View<                                                         \
+          const SCALAR_TYPE *, LAYOUT_TYPE,                                \
+          Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,                  \
+          Kokkos::MemoryTraits<Kokkos::Unmanaged | Kokkos::RandomAccess> >, \
       false, true>;
 
 #define KOKKOSSPARSE_SPTRSV_SYMBOLIC_ETI_SPEC_INST(                         \
@@ -147,6 +153,10 @@ struct SPTRSV_SYMBOLIC<ExecutionSpace, KernelHandle, RowMapType, EntriesType,
           const ORDINAL_TYPE *, LAYOUT_TYPE,                                \
           Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,                  \
           Kokkos::MemoryTraits<Kokkos::Unmanaged | Kokkos::RandomAccess> >, \
+      Kokkos::View<                                                         \
+          const SCALAR_TYPE *, LAYOUT_TYPE,                                \
+          Kokkos::Device<EXEC_SPACE_TYPE, MEM_SPACE_TYPE>,                  \
+          Kokkos::MemoryTraits<Kokkos::Unmanaged | Kokkos::RandomAccess> >, \
       false, true>;
 
 #include <KokkosSparse_sptrsv_symbolic_tpl_spec_decl.hpp>