From 2eba8ff6e259375e2098c4a6169ce83391314270 Mon Sep 17 00:00:00 2001 From: rosspeili Date: Thu, 3 Sep 2026 11:46:53 +0300 Subject: [PATCH] Avoid reduce() intermediate arrays in ParallelLinearQubitOperator (#1410) Accumulate partial matvecs into one buffer with in-place += instead of functools.reduce(numpy.add), for both the single-process and multiprocess paths. Defer forkserver setup until processes > 1 so the single-process path works on Windows. --- .../linalg/linear_qubit_operator.py | 26 +++++-- .../linalg/linear_qubit_operator_test.py | 73 ++++++++++++++++++- 2 files changed, 93 insertions(+), 6 deletions(-) diff --git a/src/openfermion/linalg/linear_qubit_operator.py b/src/openfermion/linalg/linear_qubit_operator.py index e1cd2c9d9..dd8d966be 100644 --- a/src/openfermion/linalg/linear_qubit_operator.py +++ b/src/openfermion/linalg/linear_qubit_operator.py @@ -11,7 +11,6 @@ # limitations under the License. """LinearQubitOperator is a linear operator from QubitOperator.""" -import functools import logging import multiprocessing @@ -24,6 +23,20 @@ from openfermion.config import get_available_cpu_count +def _accumulate_vectors(vectors, shape): + """Sum vectors into one array without reduce() intermediate allocations. + + ``functools.reduce(numpy.add, ...)`` builds a new full-sized array for every + partial sum. For large state vectors that creates substantial temporary + memory pressure. Accumulating with in-place ``+=`` keeps a single result + buffer instead. + """ + result = numpy.zeros(shape, dtype=complex) + for vector in vectors: + result += vector + return result + + def _bit_parity(values): """Returns the parity of the population count of each uint64 value.""" values = values ^ (values >> numpy.uint64(32)) @@ -169,7 +182,9 @@ def __init__(self, qubit_operator, n_qubits=None, options=None): self.n_qubits = n_qubits self.options = options or LinearQubitOperatorOptions() - if not ParallelLinearQubitOperator._start_method_set: + # Only required when actually spawning workers; the single-process path + # must remain usable on platforms without forkserver (e.g. Windows). + if self.options.processes > 1 and not ParallelLinearQubitOperator._start_method_set: multiprocessing.set_start_method('forkserver', force=True) ParallelLinearQubitOperator._start_method_set = True @@ -193,7 +208,9 @@ def _matvec(self, x): return numpy.zeros(x.shape) if self.options.processes <= 1: - return functools.reduce(numpy.add, (operator * x for operator in self.linear_operators)) + return _accumulate_vectors( + (operator * x for operator in self.linear_operators), x.shape + ) pool = self.options.get_pool(len(self.linear_operators)) vecs = pool.imap_unordered( @@ -203,10 +220,9 @@ def _matvec(self, x): # Consume results before join(): imap_unordered uses a bounded pipe and # workers block on write if the main process has not read them yet. try: - result = functools.reduce(numpy.add, vecs) + return _accumulate_vectors(vecs, x.shape) finally: pool.join() - return result def apply_operator(args): diff --git a/src/openfermion/linalg/linear_qubit_operator_test.py b/src/openfermion/linalg/linear_qubit_operator_test.py index 0d044cb7a..02eb69e37 100644 --- a/src/openfermion/linalg/linear_qubit_operator_test.py +++ b/src/openfermion/linalg/linear_qubit_operator_test.py @@ -22,6 +22,7 @@ LinearQubitOperator, LinearQubitOperatorOptions, ParallelLinearQubitOperator, + _accumulate_vectors, apply_operator, generate_linear_qubit_operator, ) @@ -188,6 +189,7 @@ def test_matvec_compare(self): ) +@unittest.skipIf(sys.platform == 'win32', 'forkserver multiprocessing is Unix-only') class ParallelLinearQubitOperatorTest(unittest.TestCase): """Tests for ParallelLinearQubitOperator class.""" @@ -284,6 +286,74 @@ def test_matvec_large_vector_multiprocess(self): self.assertTrue(numpy.allclose(parallel_op * vec, serial_op * vec)) +class ParallelLinearQubitOperatorSingleProcessTest(unittest.TestCase): + """Single-process ParallelLinearQubitOperator path (#1410). + + Kept separate so Windows can exercise the processes=1 accumulate path without + constructing a default multi-process operator (forkserver is Unix-only). + """ + + def setUp(self): + self.qubit_operator = QubitOperator('Z3') + QubitOperator('Y0') + QubitOperator('X1') + self.n_qubits = 4 + self.vec = numpy.array(range(2**self.n_qubits)) + self.options = LinearQubitOperatorOptions(processes=1) + self.parallel_op = ParallelLinearQubitOperator( + self.qubit_operator, self.n_qubits, options=self.options + ) + self.serial_op = LinearQubitOperator(self.qubit_operator, self.n_qubits) + + def test_matvec_matches_expected(self): + expected = numpy.array([0, -1, 2, -3, 4, -5, 6, -7, 8, -9, 10, -11, 12, -13, 14, -15]) + expected = expected + numpy.array( + [-8j, -9j, -10j, -11j, -12j, -13j, -14j, -15j, 0j, 1j, 2j, 3j, 4j, 5j, 6j, 7j] + ) + expected += numpy.array([4, 5, 6, 7, 0, 1, 2, 3, 12, 13, 14, 15, 8, 9, 10, 11]) + self.assertTrue(numpy.allclose(self.parallel_op * self.vec, expected)) + + def test_matvec_matches_serial(self): + self.assertTrue(numpy.allclose(self.parallel_op * self.vec, self.serial_op * self.vec)) + + def test_matvec_preserves_input_shape(self): + column = self.vec.reshape(-1, 1) + result = self.parallel_op * column + self.assertEqual(result.shape, column.shape) + self.assertTrue(numpy.allclose(result.reshape(-1), (self.serial_op * self.vec))) + + def test_matvec_accumulates_multiple_operator_groups(self): + """Force several groups with processes=1 to exercise in-place summing.""" + # Normal construction with processes=1 yields a single group; override the + # operator list to verify accumulation of multiple partial matvecs. + self.parallel_op.linear_operators = [ + LinearQubitOperator(QubitOperator('Z3'), self.n_qubits), + LinearQubitOperator(QubitOperator('Y0'), self.n_qubits), + LinearQubitOperator(QubitOperator('X1'), self.n_qubits), + ] + self.assertTrue(numpy.allclose(self.parallel_op * self.vec, self.serial_op * self.vec)) + + +class AccumulateVectorsTest(unittest.TestCase): + """Tests for in-place vector accumulation helper (#1410).""" + + def test_accumulate_vectors_sums_like_numpy_add(self): + vectors = [ + numpy.array([1, 2, 3], dtype=complex), + numpy.array([4, 5, 6], dtype=complex), + numpy.array([7, 8, 9], dtype=complex), + ] + expected = vectors[0] + vectors[1] + vectors[2] + self.assertTrue(numpy.allclose(_accumulate_vectors(vectors, (3,)), expected)) + + def test_accumulate_vectors_empty_iterable(self): + self.assertTrue( + numpy.allclose(_accumulate_vectors([], (4,)), numpy.zeros(4, dtype=complex)) + ) + + def test_accumulate_vectors_single_vector(self): + vector = numpy.array([1j, -2j], dtype=complex) + self.assertTrue(numpy.allclose(_accumulate_vectors([vector], (2,)), vector)) + + class UtilityFunctionTest(unittest.TestCase): """Tests for utility functions.""" @@ -308,8 +378,9 @@ def test_generate_linear_operator(self): self.assertTrue(isinstance(operator, LinearQubitOperator)) self.assertFalse(isinstance(operator, ParallelLinearQubitOperator)) + # processes=1 exercises ParallelLinearQubitOperator without forkserver. operator_again = generate_linear_qubit_operator( - qubit_operator, n_qubits, options=LinearQubitOperatorOptions(2) + qubit_operator, n_qubits, options=LinearQubitOperatorOptions(1) ) self.assertTrue(isinstance(operator_again, ParallelLinearQubitOperator)) self.assertFalse(isinstance(operator_again, LinearQubitOperator))