From 59985d99fb3d666a1e6284cc83d1b879467ac4b6 Mon Sep 17 00:00:00 2001 From: Ben Thompson Date: Wed, 29 Sep 2021 17:01:00 -0400 Subject: [PATCH 1/4] Avoid one allocation in SplitMatrix.matvec --- .../matrix/benchmark/memory_tools.py | 4 +-- src/quantcore/matrix/split_matrix.py | 28 ++++++++++++++++--- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/src/quantcore/matrix/benchmark/memory_tools.py b/src/quantcore/matrix/benchmark/memory_tools.py index aa2c2e88..3714f9e0 100644 --- a/src/quantcore/matrix/benchmark/memory_tools.py +++ b/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 diff --git a/src/quantcore/matrix/split_matrix.py b/src/quantcore/matrix/split_matrix.py index dabeae46..df6d4436 100644 --- a/src/quantcore/matrix/split_matrix.py +++ b/src/quantcore/matrix/split_matrix.py @@ -347,11 +347,31 @@ 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 the 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. + 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( From 339a3a4111e78113a4781d396054af4bd328f6f8 Mon Sep 17 00:00:00 2001 From: Ben Thompson Date: Wed, 29 Sep 2021 17:01:46 -0400 Subject: [PATCH 2/4] Avoid one allocation in SplitMatrix.matvec --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9a9b54e9..4a60bda8 100644 --- a/CHANGELOG.rst +++ b/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 ------------------ From e63fc2a8e17accdec1355c4547f2b2a748f8a92d Mon Sep 17 00:00:00 2001 From: Ben Thompson Date: Wed, 29 Sep 2021 17:03:03 -0400 Subject: [PATCH 3/4] Update split_matrix.py --- src/quantcore/matrix/split_matrix.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/quantcore/matrix/split_matrix.py b/src/quantcore/matrix/split_matrix.py index df6d4436..d0da3b3e 100644 --- a/src/quantcore/matrix/split_matrix.py +++ b/src/quantcore/matrix/split_matrix.py @@ -350,10 +350,11 @@ def matvec( # 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 the we want to + # 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. + # 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): From 13011aab51bb680002d4c17bf2731770d569c50f Mon Sep 17 00:00:00 2001 From: Ben Thompson Date: Wed, 29 Sep 2021 17:42:08 -0400 Subject: [PATCH 4/4] black fix. --- src/quantcore/matrix/split_matrix.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/quantcore/matrix/split_matrix.py b/src/quantcore/matrix/split_matrix.py index d0da3b3e..ee9cf5fc 100644 --- a/src/quantcore/matrix/split_matrix.py +++ b/src/quantcore/matrix/split_matrix.py @@ -353,7 +353,7 @@ def matvec( # 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 + # 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)