Support mixed CSR/CSC in MKL.

BenBrock · BenBrock · commit 543ed24cadd9 · 2025-03-06T15:06:21.000-08:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -6,6 +6,8 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
 set(CMAKE_CXX_FLAGS "-O3 -march=native")
 
+option(ENABLE_SANITIZERS "Enable Clang sanitizers" OFF)
+
 # Get includes, which declares the `spblas` library
 add_subdirectory(include)
 
@@ -36,6 +38,13 @@ if (LOG_LEVEL)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DLOG_LEVEL=${LOG_LEVEL}") # SPBLAS_DEBUG | SPBLAS_WARNING | SPBLAS_TRACE | SPBLAS_INFO
 endif()
 
+# Enable sanitizers
+if (ENABLE_SANITIZERS)
+    set(SANITIZER_FLAGS "-fsanitize=address,undefined")
+    target_compile_options(spblas INTERFACE ${SANITIZER_FLAGS} -g -O1 -fno-omit-frame-pointer)
+    target_link_options(spblas INTERFACE ${SANITIZER_FLAGS})
+endif()
+
 # mdspan
 FetchContent_Declare(
   mdspan
diff --git a/include/spblas/algorithms/multiply_impl.hpp b/include/spblas/algorithms/multiply_impl.hpp
@@ -212,6 +212,8 @@ std::optional<T> sparse_dot_product(A&& a, B&& b) {
     if (a_i == b_i) {
       sum += a_v * b_v;
       implicit_zero = false;
+      ++a_iter;
+      ++b_iter;
     } else if (a_i < b_i) {
       ++a_iter;
     } else {
@@ -268,6 +270,7 @@ void multiply(A&& a, B&& b, C&& c) {
       }
     }
   }
+  c_builder.finish();
   c.update(c.values(), c.rowptr(), c.colind(), c.shape(),
            c.rowptr()[c.shape()[0]]);
 }
diff --git a/include/spblas/algorithms/transpose_impl.hpp b/include/spblas/algorithms/transpose_impl.hpp
@@ -13,14 +13,15 @@ operation_info_t transpose_inspect(A&& a, B&& b) {
 
 template <matrix A, matrix B>
   requires(__detail::is_csr_view_v<A> && __detail::is_csr_view_v<B>)
-void transpose(operation_info_t& info, A&& a, B&& b) {
+void transpose(A&& a, B&& b) {
   if (__backend::shape(a)[0] != __backend::shape(b)[1] ||
       __backend::shape(a)[1] != __backend::shape(b)[0]) {
     throw std::invalid_argument(
         "transpose: matrix dimensions are incompatible.");
   }
-  if (__backend::size(a) != __backend::size(b)) {
-    throw std::invalid_argument("transpose: matrix nnz are incompatible.");
+  if (b.values().size() < __backend::size(a) ||
+      b.colind().size() < __backend::size(a)) {
+    throw std::runtime_error("transpose: Transpose ran out of memory.");
   }
   using O = tensor_offset_t<B>;
 
@@ -47,6 +48,14 @@ void transpose(operation_info_t& info, A&& a, B&& b) {
       b_rowptr[j + 1]++;
     }
   }
+
+  b.update(b.values(), b.rowptr(), b.colind(), b.shape(), a.size());
+}
+
+template <matrix A, matrix B>
+  requires(__detail::is_csr_view_v<A> && __detail::is_csr_view_v<B>)
+void transpose(operation_info_t& info, A&& a, B&& b) {
+  transpose(std::forward<A>(a), std::forward<B>(b));
 }
 
 } // namespace spblas
diff --git a/include/spblas/backend/csr_builder.hpp b/include/spblas/backend/csr_builder.hpp
@@ -26,17 +26,25 @@ class csr_builder {
 
     while (i_ < row_index) {
       view_.rowptr()[i_ + 1] = j_ptr_;
+      i_++;
     }
 
     for (auto&& [j, v] : row) {
       view_.values()[j_ptr_] = v;
       view_.colind()[j_ptr_] = j;
-      ++j_ptr_;
+      j_ptr_++;
     }
     view_.rowptr()[i_ + 1] = j_ptr_;
     i_++;
   }
 
+  void finish() {
+    while (i_ < view_.shape()[0]) {
+      view_.rowptr()[i_ + 1] = j_ptr_;
+      i_++;
+    }
+  }
+
 private:
   csr_view<T, I, O> view_;
   O j_ptr_ = 0;
diff --git a/include/spblas/vendor/onemkl_sycl/spgemm_impl.hpp b/include/spblas/vendor/onemkl_sycl/spgemm_impl.hpp
@@ -8,6 +8,7 @@
 #include <spblas/detail/operation_info_t.hpp>
 #include <spblas/detail/ranges.hpp>
 #include <spblas/detail/view_inspectors.hpp>
+#include <spblas/vendor/onemkl_sycl/detail/create_matrix_handle.hpp>
 
 //
 // Defines the following APIs for SpGEMM:
@@ -23,8 +24,9 @@
 namespace spblas {
 
 template <matrix A, matrix B, matrix C>
-  requires __detail::has_csr_base<A> && __detail::has_csr_base<B> &&
-           __detail::is_csr_view_v<C>
+  requires(__detail::has_csr_base<A> || __detail::has_csc_base<A>) &&
+          (__detail::has_csr_base<B> || __detail::has_csc_base<B>) &&
+          __detail::is_csr_view_v<C>
 operation_info_t multiply_compute(A&& a, B&& b, C&& c) {
   log_trace("");
   auto a_base = __detail::get_ultimate_base(a);
@@ -34,51 +36,40 @@ operation_info_t multiply_compute(A&& a, B&& b, C&& c) {
   using oneapi::mkl::sparse::matmat_request;
   using oneapi::mkl::sparse::matrix_view_descr;
 
-  oneapi::mkl::sparse::matmat_descr_t descr = nullptr;
-
   sycl::queue q(sycl::cpu_selector_v);
 
-  oneapi::mkl::sparse::init_matmat_descr(&descr);
-
-  oneapi::mkl::sparse::set_matmat_data(
-      descr, matrix_view_descr::general, transpose::nontrans, // view/op for A
-      matrix_view_descr::general, transpose::nontrans,        // view/op for B
-      matrix_view_descr::general);                            // view for C
-
-  oneapi::mkl::sparse::matrix_handle_t a_handle, b_handle, c_handle;
-  a_handle = b_handle = c_handle = nullptr;
-
-  oneapi::mkl::sparse::init_matrix_handle(&a_handle);
-  oneapi::mkl::sparse::init_matrix_handle(&b_handle);
-  oneapi::mkl::sparse::init_matrix_handle(&c_handle);
-
-  oneapi::mkl::sparse::set_csr_data(
-      q, a_handle, __backend::shape(a_base)[0], __backend::shape(a_base)[1],
-      oneapi::mkl::index_base::zero, a_base.rowptr().data(),
-      a_base.colind().data(), a_base.values().data())
-      .wait();
-
-  oneapi::mkl::sparse::set_csr_data(
-      q, b_handle, __backend::shape(b_base)[0], __backend::shape(b_base)[1],
-      oneapi::mkl::index_base::zero, b_base.rowptr().data(),
-      b_base.colind().data(), b_base.values().data())
-      .wait();
-
   using T = tensor_scalar_t<C>;
   using I = tensor_index_t<C>;
 
+  oneapi::mkl::sparse::matrix_handle_t a_handle =
+      __mkl::create_matrix_handle(q, a_base);
+  oneapi::mkl::sparse::matrix_handle_t b_handle =
+      __mkl::create_matrix_handle(q, b_base);
+
   I* c_rowptr;
   if (c.rowptr().size() >= __backend::shape(c)[0] + 1) {
     c_rowptr = c.rowptr().data();
   } else {
     c_rowptr = sycl::malloc_device<I>(__backend::shape(c)[0] + 1, q);
   }
 
+  oneapi::mkl::sparse::matrix_handle_t c_handle = nullptr;
+  oneapi::mkl::sparse::init_matrix_handle(&c_handle);
+
   oneapi::mkl::sparse::set_csr_data(
       q, c_handle, __backend::shape(c)[0], __backend::shape(c)[1],
       oneapi::mkl::index_base::zero, c_rowptr, (I*) nullptr, (T*) nullptr)
       .wait();
 
+  oneapi::mkl::sparse::matmat_descr_t descr = nullptr;
+  oneapi::mkl::sparse::init_matmat_descr(&descr);
+
+  oneapi::mkl::sparse::set_matmat_data(
+      descr, matrix_view_descr::general,
+      __mkl::get_transpose(a),                             // view/op for A
+      matrix_view_descr::general, __mkl::get_transpose(b), // view/op for B
+      matrix_view_descr::general);                         // view for C
+
   auto ev1 = oneapi::mkl::sparse::matmat(q, a_handle, b_handle, c_handle,
                                          matmat_request::work_estimation, descr,
                                          nullptr, nullptr, {});
@@ -113,8 +104,9 @@ operation_info_t multiply_compute(A&& a, B&& b, C&& c) {
 }
 
 template <matrix A, matrix B, matrix C>
-  requires __detail::has_csr_base<A> && __detail::has_csr_base<B> &&
-           __detail::is_csr_view_v<C>
+  requires(__detail::has_csr_base<A> || __detail::has_csc_base<A>) &&
+          (__detail::has_csr_base<B> || __detail::has_csc_base<B>) &&
+          __detail::is_csr_view_v<C>
 void multiply_fill(operation_info_t& info, A&& a, B&& b, C&& c) {
 
   log_trace("");
diff --git a/test/gtest/CMakeLists.txt b/test/gtest/CMakeLists.txt
@@ -5,6 +5,7 @@ add_executable(
   spmv_test.cpp
   spmm_test.cpp
   spgemm_test.cpp
+  spgemm_csr_csc.cpp
   add_test.cpp
   transpose_test.cpp
   triangular_solve_test.cpp
diff --git a/test/gtest/spgemm_csr_csc.cpp b/test/gtest/spgemm_csr_csc.cpp
@@ -0,0 +1,93 @@
+#include <gtest/gtest.h>
+
+#include "util.hpp"
+#include <spblas/backend/spa_accumulator.hpp>
+#include <spblas/spblas.hpp>
+
+#include <fmt/core.h>
+#include <fmt/ranges.h>
+
+TEST(MixedViews, SpGEMM_CsrCsc) {
+  using T = float;
+  using I = spblas::index_t;
+
+  for (auto&& [m, k, nnz] : util::dims) {
+    for (auto&& n : {m, k}) {
+      auto [a_values, a_rowptr, a_colind, a_shape, a_nnz] =
+          spblas::generate_csr<T, I>(m, k, nnz);
+
+      auto [b_values, b_rowptr, b_colind, b_shape, b_nnz] =
+          spblas::generate_csr<T, I>(k, n, nnz);
+
+      // We will be multiplying a times b.
+      spblas::csr_view<T, I> a(a_values, a_rowptr, a_colind, a_shape, a_nnz);
+      spblas::csr_view<T, I> b(b_values, b_rowptr, b_colind, b_shape, b_nnz);
+
+      // But we'd like the second operand to be a CSC matrix.
+      // We first transpose b.
+
+      std::vector<T> b_t_values(b.size());
+      std::vector<I> b_t_rowptr(b.shape()[1] + 1);
+      std::vector<I> b_t_colind(b.size());
+
+      spblas::csr_view<T, I> b_t(b_t_values, b_t_rowptr, b_t_colind,
+                                 {b.shape()[1], b.shape()[0]}, 0);
+
+      spblas::transpose(b, b_t);
+
+      // We then build a CSC representation of b from b_t.
+      spblas::csc_view<T, I> b_csc(b_t.values(), b_t.rowptr(), b_t.colind(),
+                                   {b_t.shape()[1], b_t.shape()[0]},
+                                   b_t.size());
+
+      // Now let's multiply a * b_csc -> c.
+
+      std::vector<I> c_rowptr(m + 1);
+      spblas::csr_view<T, I> c(nullptr, c_rowptr.data(), nullptr, {m, n}, 0);
+
+      auto info = spblas::multiply_compute(a, b_csc, c);
+
+      std::vector<T> c_values(info.result_nnz());
+      std::vector<I> c_colind(info.result_nnz());
+
+      c.update(c_values, c_rowptr, c_colind);
+
+      spblas::multiply_fill(info, a, b_csc, c);
+
+      // Now that we have c, let's compute a reference c_ref.
+      // We perform a * b -> c_ref
+
+      std::vector<I> c_ref_rowptr(m + 1);
+
+      spblas::csr_view<T, I> c_ref(nullptr, c_ref_rowptr.data(), nullptr,
+                                   {m, n}, 0);
+
+      info = spblas::multiply_compute(a, b, c_ref);
+
+      std::vector<T> c_ref_values(info.result_nnz());
+      std::vector<I> c_ref_colind(info.result_nnz());
+
+      c_ref.update(c_ref_values, c_ref_rowptr, c_ref_colind);
+
+      spblas::multiply_fill(info, a, b, c_ref);
+
+      spblas::__backend::spa_accumulator<T, I> c_row_acc(c.shape()[1]);
+
+      for (auto&& [i, c_row] : spblas::__backend::rows(c)) {
+        c_row_acc.clear();
+
+        auto&& c_ref_row = spblas::__backend::lookup_row(c_ref, i);
+
+        EXPECT_EQ(c_row.size(), c_ref_row.size());
+
+        for (auto&& [j, v] : c_row) {
+          c_row_acc[j] += v;
+        }
+
+        for (auto&& [j, v] : c_ref_row) {
+          EXPECT_EQ_(v, c_row_acc[j]);
+        }
+      }
+    }
+  }
+}

Original file line number	Diff line number	Diff line change
`@@ -212,6 +212,8 @@ std::optional<T> sparse_dot_product(A&& a, B&& b) {`
`212`	`212`	`if (a_i == b_i) {`
`213`	`213`	`sum += a_v * b_v;`
`214`	`214`	`implicit_zero = false;`
	`215`	`+ ++a_iter;`
	`216`	`+ ++b_iter;`
`215`	`217`	`} else if (a_i < b_i) {`
`216`	`218`	`++a_iter;`
`217`	`219`	`} else {`
`@@ -268,6 +270,7 @@ void multiply(A&& a, B&& b, C&& c) {`
`268`	`270`	`}`
`269`	`271`	`}`
`270`	`272`	`}`
	`273`	`+ c_builder.finish();`
`271`	`274`	`c.update(c.values(), c.rowptr(), c.colind(), c.shape(),`
`272`	`275`	`c.rowptr()[c.shape()[0]]);`
`273`	`276`	`}`