Fix memory split matvec

CHANGELOG.rst

@@ -20,6 +20,7 @@ Unreleased
 - Implemented :func:`CategoricalMatrix.__rmul__`
 - Enable indexing the rows of a ``CategoricalMatrix``. Previously :func:`CategoricalMatrix.__getitem__` only supported column indexing.
 - Allow creating a ``SplitMatrix`` from a list of any ``MatrixBase`` objects including another ``SplitMatrix``.
+- Reduced memory usage in :meth:`quantcore.matrix.SplitMatrix.matvec`.
 
 2.0.3 - 2021-07-15
 ------------------

src/quantcore/matrix/benchmark/memory_tools.py

@@ -40,7 +40,7 @@ def track_peak_mem(f, *args, **kwargs):
         f(*args, **kwargs)
     for s in mp.snapshots:
         top_stats = s.statistics("lineno")
-        print("[ Top 2 ]")
-        for stat in top_stats[:2]:
+        print("[ Top 4 ]")
+        for stat in top_stats[:4]:
             print(stat)
     return mp.peak_usage

src/quantcore/matrix/split_matrix.py

@@ -347,11 +347,32 @@ def matvec(
 
         out_shape = [self.shape[0]] + ([] if v.ndim == 1 else list(v.shape[1:]))
         out_dtype = np.result_type(self.dtype, v.dtype)
-        out = _prepare_out_array(out, out_shape, out_dtype)
 
-        for sub_cols, idx, mat in zip(subset_cols, self.indices, self.matrices):
-            one = v[idx, ...]
-            mat.matvec(one, sub_cols, out=out)
+        # If there is a dense matrix in the list of matrices, we want to
+        # multiply that one first for memory use reasons. This is because numpy
+        # doesn't provide a blas-like mechanism for specifying that we want to
+        # add the result of the matrix-vector product into an existing array.
+        # So, we simply use the output of the first dense matrix-vector product
+        # as the target for storing the final output. This reduces the number
+        # of output arrays allocated from 2 to 1.
+        is_matrix_dense = [isinstance(m, DenseMatrix) for m in self.matrices]
+        dense_matrix_idx = np.argmax(is_matrix_dense)
+        if np.any(is_matrix_dense):
+            sub_cols = subset_cols[dense_matrix_idx]
+            idx = self.indices[dense_matrix_idx]
+            mat = self.matrices[dense_matrix_idx]
+            in_vec = v[idx, ...]
+            out = np.asarray(mat.matvec(in_vec, sub_cols, out), dtype=out_dtype)
+        else:
+            out = _prepare_out_array(out, out_shape, out_dtype)
+
+        for i, (sub_cols, idx, mat) in enumerate(
+            zip(subset_cols, self.indices, self.matrices)
+        ):
+            if i == dense_matrix_idx:
+                continue
+            in_vec = v[idx, ...]
+            mat.matvec(in_vec, sub_cols, out=out)
         return out
 
     def transpose_matvec(