Modifications for rocsparse implementation for SpGEMM

- Adding Jacobi SpGEMM via rocsparse - Enabling rocsparse testing for SpGEMM - Enabling complex datatypes with SpGEMM - Adding option for non-int size types (e.g. size_t and int64_t) for SpGEMM
kokkos · Nov 15, 2022 · 43dd4a7 · 43dd4a7
1 parent f3556c5
commit 43dd4a7
Show file tree

Hide file tree

Showing 8 changed files with 842 additions and 241 deletions.
diff --git a/sparse/impl/KokkosSparse_par_ilut_numeric_impl.hpp b/sparse/impl/KokkosSparse_par_ilut_numeric_impl.hpp
@@ -953,8 +953,8 @@ struct IlutWrap {
         const index_t l_nnz = L_new_values.extent(0);
         const index_t u_nnz = U_new_values.extent(0);
 
-        const auto l_filter_rank = std::max(0, l_nnz - l_nnz_limit - 1);
-        const auto u_filter_rank = std::max(0, u_nnz - u_nnz_limit - 1);
+        const auto l_filter_rank = std::max(index_t(0), l_nnz - l_nnz_limit - 1);
+        const auto u_filter_rank = std::max(index_t(0), u_nnz - u_nnz_limit - 1);
 
         const auto l_threshold =
             threshold_select(L_new_values, l_filter_rank, V_copy);

diff --git a/sparse/impl/KokkosSparse_spgemm_jacobi_spec.hpp b/sparse/impl/KokkosSparse_spgemm_jacobi_spec.hpp
@@ -57,6 +57,10 @@
 #include "KokkosSparse_spgemm_jacobi_seq_impl.hpp"
 #endif
 
+#if defined(KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE)
+#include "KokkosSparse_spgemm_rocSPARSE_impl.hpp"
+#endif
+
 namespace KokkosSparse {
 namespace Impl {
 // Specialization struct which defines whether a specialization exists
@@ -194,6 +198,20 @@ struct SPGEMM_JACOBI<KernelHandle, a_size_view_t_, a_lno_view_t,
       spgemm_jacobi_seq(handle, m, n, k, row_mapA, entriesA, valuesA,
                         transposeA, row_mapB, entriesB, valuesB, transposeB,
                         row_mapC, entriesC, valuesC, omega, dinv);
+    } else if (sh->get_algorithm_type() == SPGEMM_ROCSPARSE) {
+#if defined(KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE)
+      jacobi_spgemm_numeric_rocsparse<spgemmHandleType,
+          a_size_view_t_,a_lno_view_t,a_scalar_view_t,
+          c_lno_view_t,c_scalar_view_t,dinv_scalar_view_t>(
+            sh, m, n, k,
+            row_mapA, entriesA, valuesA, transposeA,
+            row_mapB, entriesB, valuesB, transposeB,
+            omega, dinv,
+            row_mapC, entriesC, valuesC);
+#else
+      throw std::runtime_error(
+          "Requiring (jacobi) SPGEMM_ROCSPARSE but TPL_ROCSPARSE was not enabled!");
+#endif
     } else {
       KokkosSPGEMM<KernelHandle, a_size_view_t_, a_lno_view_t, a_scalar_view_t,
                    b_size_view_t_, b_lno_view_t, b_scalar_view_t>
@@ -292,7 +310,7 @@ struct SPGEMM_JACOBI<KernelHandle, a_size_view_t_, a_lno_view_t,
                    Kokkos::MemoryTraits<Kokkos::Unmanaged> >,             \
       false, true>;
 
-#include <KokkosSparse_spgemm_tpl_spec_decl.hpp>
+// #include <KokkosSparse_spgemm_tpl_spec_decl.hpp>
 #include <generated_specializations_hpp/KokkosSparse_spgemm_jacobi_eti_spec_decl.hpp>
 
 #endif
diff --git a/sparse/impl/KokkosSparse_spgemm_numeric_spec.hpp b/sparse/impl/KokkosSparse_spgemm_numeric_spec.hpp
@@ -223,16 +223,6 @@ struct SPGEMM_NUMERIC<
 #else
         throw std::runtime_error(
             "Requiring SPGEMM_CUSPARSE but TPL_CUSPARSE was not enabled!");
-#endif
-        break;
-      case SPGEMM_ROCSPARSE:
-#if defined(KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE)
-        rocsparse_spgemm_numeric<spgemmHandleType>(
-            sh, m, n, k, row_mapA, entriesA, valuesA, transposeA, row_mapB,
-            entriesB, valuesB, transposeB, row_mapC, entriesC, valuesC);
-#else
-        throw std::runtime_error(
-            "Requiring SPGEMM_ROCSPARSE but TPL_ROCSPARSE was not enabled!");
 #endif
         break;
       case SPGEMM_CUSP:
@@ -243,6 +233,18 @@ struct SPGEMM_NUMERIC<
                                     transposeA, row_mapB, entriesB, valuesB,
                                     transposeB, row_mapC, entriesC, valuesC);
         break;
+      case SPGEMM_ROCSPARSE:
+#if defined(KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE)
+        spgemm_numeric_rocsparse<spgemmHandleType, a_size_view_t_, a_lno_view_t,a_scalar_view_t,c_lno_view_t,c_scalar_view_t>(
+            sh, m, n, k,
+            row_mapA, entriesA, valuesA, transposeA,
+            row_mapB, entriesB, valuesB, transposeB,
+            row_mapC, entriesC, valuesC);
+#else
+        throw std::runtime_error("Requested rocSPARSE for SpGEMM, but it is not available. Please recompile with rocsparse enabled.");
+#endif
+        break;
+
       case SPGEMM_MKL:
 #ifdef KOKKOSKERNELS_ENABLE_TPL_MKL
         mkl_numeric(sh, m, n, k, row_mapA, entriesA, valuesA, transposeA,