Use csr_matvec instead of mkl for sparse matvec (#82)

* use ben algo

* update changelog

* remove sparse_dot_mkl in env

* matvec dim check

* bug fix

* changelog

* rearrange

* Update CHANGELOG.rst

Co-authored-by: Jan Tilly <[email protected]>

* Update CHANGELOG.rst

Co-authored-by: Jan Tilly <[email protected]>

* remove dot_product_mkl from conda recipe

Co-authored-by: Jan Tilly <[email protected]>

CHANGELOG.rst

@@ -7,6 +7,13 @@
 Changelog
 =========
 
+2.0.1 - 2021-06-15
+------------------
+
+**Other changes**:
+
+We removed the dependency on ``sparse_dot_mkl``. We now use ``scipy.sparse.csr_matvec`` instead of ``sparse_dot_mkl.dot_product_mkl`` on all platforms, because the former suffered from poor performance, especially on narrow problems.
+
 2.0.0 - 2021-06-10
 ------------------
 

conda.recipe/meta.yaml

@@ -44,7 +44,6 @@ requirements:
     - pandas
     - scikit-learn >=0.23
     - scipy
-    - sparse_dot_mkl >=0.4.0  # [not (osx and arm64)]
 
 test:
   requires:

environment.yml

@@ -19,7 +19,6 @@ dependencies:
   - pytest
   - pytest-xdist
   - setuptools_scm
-  - sparse_dot_mkl>=0.6.0
   - xsimd
   - matplotlib
   - seaborn

src/quantcore/matrix/benchmark/main.py

@@ -35,16 +35,14 @@ def _transpose_matvec(
     mat: Union[mx.MatrixBase, np.ndarray, sps.csc_matrix], vec: np.ndarray
 ):
     if isinstance(mat, (mx.MatrixBase, mx.StandardizedMatrix)):
-        out = np.zeros(mat.shape[1])
-        return mat.transpose_matvec(vec, out=out)
+        return mat.transpose_matvec(vec, out=None)
     else:
         return mat.T.dot(vec)
 
 
 def _matvec(mat, vec: np.ndarray) -> np.ndarray:
     if isinstance(mat, (mx.MatrixBase, mx.StandardizedMatrix)):
-        out = np.zeros(mat.shape[0])
-        return mat.matvec(vec, out=out)
+        return mat.matvec(vec, out=None)
     else:
         return mat.dot(vec)
 

src/quantcore/matrix/sparse_matrix.py

@@ -1,4 +1,3 @@
-import platform
 from typing import List, Optional, Union
 
 import numpy as np
@@ -19,20 +18,6 @@
     setup_restrictions,
 )
 
-if platform.machine() in {"arm64", "aarch64"}:
-
-    def _dot_product_maybe_mkl(matrix_a, matrix_b, out):
-        mult = matrix_a.dot(matrix_b)
-        if out is None:
-            return mult
-        out[:] += mult
-        return out
-
-
-else:
-    # Intel MKL obviously is only available on Intel-based architectures
-    from sparse_dot_mkl import dot_product_mkl as _dot_product_maybe_mkl  # type: ignore
-
 
 class SparseMatrix(sps.csc_matrix, MatrixBase):
     """A scipy.sparse csc matrix subclass."""
@@ -129,50 +114,38 @@ def _matvec_helper(
         out: Optional[np.ndarray],
         transpose: bool,
     ):
-        X = self.T if transpose else self
+        match_dim = 0 if transpose else 1
+        vec_shape = np.array(vec).shape
+        if self.shape[match_dim] != vec_shape[0]:
+            raise ValueError(
+                f"shapes {self.shape} and {vec_shape} not aligned:"
+                f"{self.shape[match_dim]} (dim {match_dim}) != {vec_shape[0]} (dim 0)"
+            )
         matrix_matvec = lambda x, v: sps.csc_matrix.dot(x, v)
         if transpose:
             matrix_matvec = lambda x, v: sps.csr_matrix.dot(x.T, v)
 
         vec = np.asarray(vec)
-
-        # NOTE: We assume that rows and cols are unique
-        unrestricted_rows = rows is None or len(rows) == self.shape[0]
-        unrestricted_cols = cols is None or len(cols) == self.shape[1]
-        if unrestricted_rows and unrestricted_cols:
-            if vec.ndim == 1:
-                return _dot_product_maybe_mkl(X, vec, out=out)
-            elif vec.ndim == 2 and vec.shape[1] == 1:
-                out_arr = None if out is None else out[:, 0]
-                return _dot_product_maybe_mkl(X, vec[:, 0], out=out_arr)[:, None]
-            res = matrix_matvec(self, vec)
-            if out is None:
-                return res
-            out += res
-            return out
+        rows, cols = setup_restrictions(self.shape, rows, cols, dtype=self.idx_dtype)
+        if transpose:
+            fast_fnc = lambda v: csc_rmatvec(self, v, rows, cols)
         else:
-            rows, cols = setup_restrictions(
-                self.shape, rows, cols, dtype=self.idx_dtype
+            fast_fnc = lambda v: csr_matvec(self.x_csr, v, rows, cols)
+        if vec.ndim == 1:
+            res = fast_fnc(vec)
+        elif vec.ndim == 2 and vec.shape[1] == 1:
+            res = fast_fnc(vec[:, 0])[:, None]
+        else:
+            res = matrix_matvec(
+                self[np.ix_(rows, cols)], vec[rows] if transpose else vec[cols]
             )
-            if transpose:
-                fast_fnc = lambda v: csc_rmatvec(self, v, rows, cols)
-            else:
-                fast_fnc = lambda v: csr_matvec(self.x_csr, v, rows, cols)
-            if vec.ndim == 1:
-                res = fast_fnc(vec)
-            elif vec.ndim == 2 and vec.shape[1] == 1:
-                res = fast_fnc(vec[:, 0])[:, None]
-            else:
-                res = matrix_matvec(
-                    self[np.ix_(rows, cols)], vec[rows] if transpose else vec[cols]
-                )
-            if out is None:
-                return res
-            if transpose:
-                out[cols] += res
-            else:
-                out[rows] += res
-            return out
+        if out is None:
+            return res
+        if transpose:
+            out[cols] += res
+        else:
+            out[rows] += res
+        return out
 
     def matvec(self, vec, cols: np.ndarray = None, out: np.ndarray = None):
         """Perform self[:, cols] @ other[cols]."""