Adding more support for SpGEMM via rocsparse

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);

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,6 @@ 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 <generated_specializations_hpp/KokkosSparse_spgemm_jacobi_eti_spec_decl.hpp>
 
 #endif

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,

sparse/impl/KokkosSparse_spgemm_symbolic_spec.hpp

@@ -161,11 +161,13 @@ struct SPGEMM_SYMBOLIC<KernelHandle, a_size_view_t_, a_lno_view_t,
         break;
       case SPGEMM_ROCSPARSE:
 #if defined(KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE)
-        rocsparse_spgemm_symbolic<spgemmHandleType, a_size_view_t_,
-                                  a_lno_view_t, b_size_view_t_, b_lno_view_t,
-                                  c_size_view_t_>(
-            sh, m, n, k, row_mapA, entriesA, transposeA, row_mapB, entriesB,
-            transposeB, row_mapC);
+        spgemm_symbolic_rocsparse<spgemmHandleType,
+          typename KernelHandle::nnz_scalar_t, a_size_view_t_, 
+          a_lno_view_t,c_size_view_t_>(
+            sh, m, n, k,
+            row_mapA, entriesA, transposeA,
+            row_mapB, entriesB, transposeB,
+            row_mapC);
 #else
         throw std::runtime_error(
             "Requiring SPGEMM_ROCSPARSE but TPL_ROCSPARSE was not enabled!");

sparse/src/KokkosSparse_spgemm_handle.hpp

@@ -149,44 +149,83 @@ class SPGEMMHandle {
       nnz_lno_persistent_work_host_view_t;  // Host view type
 
 #ifdef KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE
-  struct rocSparseSpgemmHandleType {
-    KokkosKernels::Experimental::Controls
-        kkControls;  // give a singleton rocsparse handle
-    rocsparse_handle rocsparseHandle;
-    rocsparse_operation opA, opB;
-    rocsparse_mat_descr descr_A, descr_B, descr_C, descr_D;
-    rocsparse_mat_info info_C;
-    size_t bufferSize;
-    void *buffer;
-    bool C_populated;
 
-    rocSparseSpgemmHandleType(bool transposeA, bool transposeB) {
-      opA =
-          transposeA ? rocsparse_operation_transpose : rocsparse_operation_none;
-      opB =
-          transposeB ? rocsparse_operation_transpose : rocsparse_operation_none;
+  struct rocSparseSpgemmHandleType
+  {
 
-      bufferSize  = 0;
-      buffer      = NULL;
-      C_populated = false;
+    /// Constructor
+    rocSparseSpgemmHandleType(bool transposeA, bool transposeB):
+      opA(transposeA ? rocsparse_operation_transpose : rocsparse_operation_none),
+      opB(transposeB ? rocsparse_operation_transpose : rocsparse_operation_none),
+      enable_jacobi(false),
+      C_populated(false)
+    {
+      // Get singleton handle
+      handle = kkControls.getRocsparseHandle();
+
+      // Create matrix descriptors
       KOKKOS_ROCSPARSE_SAFE_CALL_IMPL(rocsparse_create_mat_descr(&descr_A));
       KOKKOS_ROCSPARSE_SAFE_CALL_IMPL(rocsparse_create_mat_descr(&descr_B));
       KOKKOS_ROCSPARSE_SAFE_CALL_IMPL(rocsparse_create_mat_descr(&descr_C));
       KOKKOS_ROCSPARSE_SAFE_CALL_IMPL(rocsparse_create_mat_descr(&descr_D));
+
+      // Create C info object
       KOKKOS_ROCSPARSE_SAFE_CALL_IMPL(rocsparse_create_mat_info(&info_C));
-      rocsparseHandle = kkControls.getRocsparseHandle();
+
     }
 
-    ~rocSparseSpgemmHandleType() {
-      rocsparse_destroy_mat_info(info_C);
-      rocsparse_destroy_mat_descr(descr_A);
-      rocsparse_destroy_mat_descr(descr_B);
-      rocsparse_destroy_mat_descr(descr_C);
-      rocsparse_destroy_mat_descr(descr_D);
+    // Destructor
+    ~rocSparseSpgemmHandleType()
+    {
+
+      (void) rocsparse_destroy_mat_descr(descr_A);
+      (void) rocsparse_destroy_mat_descr(descr_B);
+      (void) rocsparse_destroy_mat_descr(descr_C);
+      (void) rocsparse_destroy_mat_descr(descr_D);
+      (void) rocsparse_destroy_mat_info(info_C);
     }
+
+    // rocsparse handle wrappers
+    KokkosKernels::Experimental::Controls kkControls;  // give a singleton rocsparse handle
+    rocsparse_handle handle;
+    rocsparse_operation opA, opB;
+
+    // Matrix descriptors
+    rocsparse_mat_descr descr_A;
+    rocsparse_mat_descr descr_B;
+    rocsparse_mat_descr descr_C;
+    rocsparse_mat_descr descr_D;
+    rocsparse_mat_info info_C;
+
+    // Scratch buffer allocation
+    using buffer_t = Kokkos::View<char*, 
+      Kokkos::Device<Kokkos::Experimental::HIP, Kokkos::Experimental::HIPSpace>>;
+    buffer_t buffer;
+
+    // Currently SpGEMM in rocSPARSE only supports the int datatype
+    // We use these arrays to transform back and forth as needed
+    using rocsparse_index_array_t = Kokkos::View<rocsparse_int*, 
+      Kokkos::Device<Kokkos::Experimental::HIP, Kokkos::Experimental::HIPSpace>>;
+    rocsparse_index_array_t csr_row_ptr_A;
+    rocsparse_index_array_t csr_row_ptr_B;
+    rocsparse_index_array_t csr_row_ptr_C;
+    rocsparse_index_array_t csr_row_ptr_D;
+    rocsparse_index_array_t csr_col_ind_A;
+    rocsparse_index_array_t csr_col_ind_B;
+    rocsparse_index_array_t csr_col_ind_C;
+    rocsparse_index_array_t csr_col_ind_D;
+
+    // (Jacobi only) Used to store D^{-1} * A entries for use with rocsparse csrgemm
+    scalar_temp_work_view_t csr_values_invDA;
+
+    // Used to trigger Jacobi SpGEMM via rocsparse
+    bool enable_jacobi;
+
+    // Used to ensure C column indexes are not recomputed in the spgemm_numeric call
+    bool C_populated;
   };
 
-#endif
+#endif // KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE
 
 #ifdef KOKKOSKERNELS_ENABLE_TPL_CUSPARSE
 #if (CUDA_VERSION >= 11000)

sparse/unit_test/Test_Sparse_spgemm.hpp

@@ -295,6 +295,10 @@ void test_spgemm(lno_t m, lno_t k, lno_t n, size_type nnz, lno_t bandwidth,
   algorithms.push_back(SPGEMM_MKL);
 #endif
 
+#if defined(KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE) && defined(TEST_HIP_SPARSE_CPP)
+  algorithms.push_back(SPGEMM_ROCSPARSE);
+#endif
+
   for (auto spgemm_algorithm : algorithms) {
     const uint64_t max_integer = Kokkos::ArithTraits<int>::max();
     std::string algo           = "UNKNOWN";
@@ -309,7 +313,13 @@ void test_spgemm(lno_t m, lno_t k, lno_t n, size_type nnz, lno_t bandwidth,
         is_expected_to_fail = true;
 #endif
         break;
-
+      case SPGEMM_ROCSPARSE: algo = "SPGEMM_ROCSPARSE"; 
+      // The rocsparse implementation internally casts size types to int
+      if (A.values.extent(0) > max_integer) {
+        is_expected_to_fail = true;
+      }
+
+      break;
       case SPGEMM_MKL: algo = "SPGEMM_MKL";
 #ifdef KOKKOSKERNELS_ENABLE_TPL_MKL
         if (!KokkosSparse::Impl::mkl_is_supported_value_type<scalar_t>::value) {

sparse/unit_test/Test_Sparse_spgemm_jacobi.hpp

@@ -100,6 +100,11 @@ int run_spgemm_jacobi(crsMat_t input_mat, crsMat_t input_mat2,
 
   kh.create_spgemm_handle(spgemm_algorithm);
 
+#if defined(KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE) && defined(TEST_HIP_SPARSE_CPP)
+  kh.get_spgemm_handle()->create_rocsparse_spgemm_handle(false,false);
+  kh.get_spgemm_handle()->get_rocsparse_spgemm_handle()->enable_jacobi=true;
+#endif
+
   const size_t num_rows_1 = input_mat.numRows();
   const size_t num_rows_2 = input_mat2.numRows();
   const size_t num_cols_2 = input_mat2.numCols();
@@ -256,15 +261,20 @@ void test_spgemm_jacobi(lno_t numRows, size_type nnz, lno_t bandwidth,
   run_spgemm_jacobi<crsMat_t, device, scalar_t, view_t>(
       input_mat, input_mat, omega, dinv, SPGEMM_SERIAL, output_mat2);
 
-  SPGEMMAlgorithm spgemm_algorithm =
-      SPGEMM_KK_MEMORY;  // should we test other SpGEMM algorithms as well?
+  std::vector<SPGEMMAlgorithm> spgemm_algorithms;
+  spgemm_algorithms.push_back(SPGEMM_KK_MEMORY);
+#if defined(KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE) && defined(TEST_HIP_SPARSE_CPP)
+  spgemm_algorithms.push_back(SPGEMM_ROCSPARSE);
+#endif
 
-  crsMat_t output_mat;
+  for(const auto & spgemm_algorithm : spgemm_algorithms){
+    crsMat_t output_mat;
 
-  run_spgemm_jacobi<crsMat_t, device>(input_mat, input_mat, omega, dinv,
-                                      spgemm_algorithm, output_mat);
-  bool is_identical = is_same_mat<crsMat_t, device>(output_mat, output_mat2);
-  EXPECT_TRUE(is_identical);
+    run_spgemm_jacobi<crsMat_t, device>(input_mat, input_mat, omega, dinv,
+                                        spgemm_algorithm, output_mat);
+    bool is_identical = is_same_mat<crsMat_t, device>(output_mat, output_mat2);
+    EXPECT_TRUE(is_identical);
+  }
 }
 
 #define KOKKOSKERNELS_EXECUTE_TEST(SCALAR, ORDINAL, OFFSET, DEVICE)            \