Numba-jit statevector backend#518
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #518 +/- ##
==========================================
- Coverage 90.58% 88.91% -1.68%
==========================================
Files 49 49
Lines 7481 7685 +204
==========================================
+ Hits 6777 6833 +56
- Misses 704 852 +148 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| random_op3 = rand_unit(2**3, rng) | ||
| random_op3_exp = rand_unit(2**3, rng) |
There was a problem hiding this comment.
Could you explain this change? Could we systematically replace random_op(n, rng) by rand_unit(2**n, rng) and remove random_op definition?
There was a problem hiding this comment.
In the old API, Statevec.expectation_value was normalizing the statevector before computing the norm. Now, for efficiency, we don't normalize the statevector, because we assume that the statevector always has norm 1. This is guaranteed upon instantianting a new Statevec object, and requires that Statevec.evolve only takes unitary operators, otherwise the norm is not conserved (same applies to Statevec.evolve_single, but not relevant here). I think this restriction is reasonable, as we expect gates to be unitary operations.
In general, Statevec.expectation_value doesn't need to be restricted to unitary operators, so I think we can keep random_op for the other instances.
| self._psi = data._psi.copy() | ||
| self._max_qubits = data.max_qubits | ||
| self._nqubit = data.nqubit | ||
|
|
||
| if max_qubits is not None: | ||
| if max_qubits < data.max_qubits: | ||
| raise ValueError( | ||
| f"`max_qubits` can't be smaller than the capacity of input state: {max_qubits} < {data.max_qubits}." | ||
| ) | ||
| self.ensure_capacity(max_qubits) |
There was a problem hiding this comment.
The implemented behavior is:
- Copy
data._psias‑is. - Require the user‑supplied
max_qubitsto be greater thandata.max_qubits.
I propose the following, which I think is both more useful and more intuitive:
- Require the user‑supplied
max_qubitsto be greater thandata.nqubit. - Truncate
data._psiat the user‑suppliedmax_qubitswhen necessary.
| offset = (1 << required_qubits) - len(self._psi) | ||
| self._psi = np.concatenate([self._psi, np.empty(offset, dtype=self._psi.dtype)]) |
There was a problem hiding this comment.
I've only performed micro-benchmarks but it seems much more efficient:
| offset = (1 << required_qubits) - len(self._psi) | |
| self._psi = np.concatenate([self._psi, np.empty(offset, dtype=self._psi.dtype)]) | |
| new_psi = np.empty(1 << required_qubits, dtype=self._psi.dtype) | |
| new_psi[:len(self._psi)] = self._psi | |
| self._psi = new_psi |
There was a problem hiding this comment.
Done in 696e2a2
Makes sense, we allocate once instead of twice.
2 participants
This PR replaces the statevector backend with a new version orders of magnitude more efficient. It is largely inspired from Ref. [1].
Key features:
Functions modifying the quantum state (exponentially expensive) are jit compiled with numba.
Modifications of the quantum state are done in-place so data is not unnecessarily copied during the simulation. To that purpose,
StatevectorBackendhas a new class method constructorwith_capacitywhich preallocates a given space in memory. Typically, we pass it a parametern = Pattern.max_space()to allocate a complex-type array with size2**n. ThePatternSimualatorclass handles this internally, so the API for pattern or circuit simulation remains the same.Adaptive parallelization: jit-compiled functions are parallelized for qubit counts larger than
NUM_QUBIT_PARALLEL. This compile constant (empirically determined to be 15), is the number above which the multi-thread overhead does not compensate the speed gains.In addition:
graphix-symbolicplugin. See accompanying PR: Add symbolic backends graphix-symbolic#9DenseStateare modified to follow a consistent naming. This effectively changes the public API.test_statevec.pyandtest_statevectorbackend.pyare unified in a single file.Discussion
Below I show the profiling of simulating a QFT circuit on 23 qubits (from the Munich Benchmark suite) . With the current transpiler, this pattern has
max_space=24and the following command count:{'N': 2523, 'E': 3023, 'M': 2523, 'Z': 23, 'X': 23}:Execution time is dominated by the measurement process (in particular,
remove_qubitsubcall), so further optimization efforts should focus on this part of the simulation pipeline.Current implementation is probably as good as we can do with numba. I recommend reading this very insighful thread: https://stackoverflow.com/questions/79948374/improving-efficiency-of-numba-jit-function (credits to Jérôme Richard @zephyr111).
If we want to push this further, it may be worth to write specialized code for measurements on the XY, XZ and YZ planes. Currently, measurement needs four passes over
psi(see methodgraphix.sim.base_backend.DenseStateBackend.measure):expectation_single).evolve_single).remove_qubit).remove_qubit).Step 1) is not necessary if we use the constant branch selector (this is legit for deterministic patterns, and btw, the pattern above runs in 78.9s with the 0-branch selector).
I suspect it is possible to merge steps 2) and 3) in a single pass which may give a ~30% speed improvement.
P.S. Comments by @thierry-martinez in the preliminary discussion have been addressed.
References
[1] McGuffin, M. J., Robert J-M., and Ikeda K. "How to Write a Simulator for Quantum Circuits from Scratch: A Tutorial.", 2025 (arXiv:2506.08142).
UPDATE (01/06/2026)
In commit 3ddc462 I reimplemented the method
StatevectorBackend.measureto mergeevolve_singleandremove_qubitinto a single callproject_qubitas discussed above. Effectively, we are fusing two loops overpsiwhich is advantageous in memory-bound programs. New implementation of qubit measurement is about a 23% faster:In the new implementation of the qubit-removal (step 3), we compute the squared norm of the two subvectors in the same loop. This is not the most efficient choice, since most of the times we only need the norm of one subvector. However, I observed that the difference is negligible in the example with 24 qubits, and calculating it separately leads to much more code repetition. Note that in both cases, we have to read all elements of
psito apply the projector operator which is probably more costly than doing the arithmetic operation to compute the squared norm.TODO
Update CHANGELOG