Reimplement PauliString with a symplectic bit representation#65
Open
JaipreethTiruvaipati wants to merge 1 commit into
Open
Reimplement PauliString with a symplectic bit representation#65JaipreethTiruvaipati wants to merge 1 commit into
JaipreethTiruvaipati wants to merge 1 commit into
Conversation
Replace the dict-based PauliString with two integer bitmasks (the symplectic X/Z representation), turning multiplication, (anti)commutation, collapsing and cover-vector encoding into pure integer bit operations and removing the repeated stim.PauliString <-> dict conversions. - pauli.py: symplectic masks + __slots__ + precomputed hash; O(1) exact and commuting cover vectors; bit-packed stim conversions; a CollapsingOperators value type that stores a boundary's single-qubit collapse as combined masks. - boundary.py: cache after_collapse; represent collapsing operations by their CollapsingOperators instead of a frozenset of PauliString objects. - cover.py: encode the GF(2) cover system over whole strings. - flow.py: share per-fragment collapse masks; find the measurements/resets a stabilizer overlaps with a bitwise AND instead of scanning every operator. - match.py: replace the loop-sanity deepcopy with a shallow flow clone. The annotated detectors are unchanged; on tqec surface-code memory circuits this is ~7x faster than the dict baseline and the gap grows with code size. Fixes tqec#47
Contributor
|
This is quite a chunky PR and #63 exists and has already been through several rounds of review. Because the hackathon is close to an end, we might not have the time to iterate on this PR. I'll try to provide you feedback on it anyway, but know that there is a high chance that #63 is merged instead, because I am reasonably certain it is correct after 2 rounds of review. In that case, my review will likely focus on the differences between this PR and #63 to find where the additional performance is gained in this PR. |
Contributor
Author
|
Sure, @nelimee... Will keep it open... Lets let me know what specific optimizations you might need.. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What I changed
I went after #47 by profiling first, and the issue is right: on the surface-code circuits almost all of the time really is spent inside
PauliString. So I rebuilt it on the symplectic representation (two integer bitmasks) and then followed the profile intoboundary,cover,flowandmatchuntil the per-string work stopped showing up at the top.Concretely:
pauli.py- a Pauli string is now two integers,xsandzs. Multiplication is an XOR, (anti)commutation is a popcount, collapse is a mask clear, so the hot operations are plain integer arithmetic and the constantstim.PauliString⇄dictround-trips go away. I also added__slots__and a cached hash, bit-packed NumPy for the stim conversions, cover vectors built straight from the masks, and a smallCollapsingOperatorsvalue type that holds a boundary's single-qubit collapse as combined masks.boundary.py-after_collapseis cached, and the collapsing operations are kept as theirCollapsingOperatorsinstead of afrozensetof one-qubitPauliStrings. That turns the anti-commutation check and the collapse into single bit operations rather than a loop over operators, which was taking a surprisingly large share of the runtime.cover.py- the GF(2) cover is solved over whole-string symplectic vectors.flow.py- each fragment's collapse masks are computed once, and "which measurements/resets does this tabilizer touch" becomes a bitwise AND plus a lookup instead of a scan over every operator.match.py- the loop sanity-check was doing adeepcopy; since the boundary stabilizers are immutable I swapped it for a shallow clone of the flow lists.The public API of
PauliStringis unchanged, and the detectors that come out are identical tomaineverywhere I checked.Numbers
Benchmark environment:
6782ed5, this branch5685fdaThe k = 1..9
tqecmemory(Basis.Z)benchmark from the issue, median ofseveral runs:
The k = 1 case is actually a touch slower - at 25 qubits the bitmask setup costs bit more than it saves - but from k = 2 the gap opens up fast, reaching about 7.7× at k = 9, and the ratio keeps growing with the code distance. One note on
reading those multipliers: they are measured against
mainon the hardware above, and the multiplier gets larger the slower that baseline machine is, so the absolute times are the more meaningful figure. Detector counts and the fullannotated output match
mainat every k.With this in, the profile is dominated by stim's tableau construction and
_try_merge_anticommuting_flows_inplacerather thanPauliString, which is what the issue predicted would happen. I'm happy to take that on next in a separate PR if it would be useful.Checks
pytestsuite passes, including the end-to-end round-trip over thetest_filescorpus. I added tests for the symplectic operations, the cover vectors,CollapsingOperators, theafter_collapsecache and the shallow flow clone. CI covers the 3.10–3.14 matrix.ruff check,ruff format --check,mypy src/tqecdand the fullpre-commitset are green.On AI use
I used an AI assistant while working on this - for profiling, parts of the implementation, the benchmark harness, and a first draft of this description. I read through the whole diff, ran and checked the tests, types, lint and benchmarks myself, and I understand the design: in particular why a boundary's single-qubit collapse can be reduced to combined
(X, Z)masks, and why((xs & cz) ^ (zs & cx)) != 0is the same thing as "at least one collapsing operator anti-commutes with the stabilizer". I looked at how the problem has been approached before and wrote my own version. Glad to walk through any part of it.Fixes #47