Repository navigation
Run strategy="active_statevector" in numba kernels, without JAX - #199
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e394e8404a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _, order = np.unique(packed[first], axis=0, return_index=True) | ||
| rank = np.empty_like(order) | ||
| rank[order] = np.arange(len(order)) | ||
| return matrix[:, first[order]], rank[column_ids] |
There was a problem hiding this comment.
Normalize columns before returning packed representatives
When matrix contains noncanonical nonzero values, which the docstring explicitly says count as 1, np.packbits coalesces them with literal ones but this return keeps the original representative. For example, columns containing 2 and 1 receive one ID, and if the 2 column comes first, ChannelSampler stores 2 in signature_matrix; the later matrix multiplication modulo 2 then treats both columns as zero, so sampled errors stop flipping their outputs. Normalize the matrix to boolean/0-or-1 before both packing and returning it (or reject values outside {0, 1}) to avoid order-dependent noise samples.
Useful? React with 👍 / 👎.
☂️ Code Coverage
Overall Coverage
New Files
Modified Files
|
The survivors of each chunk were appended to a Python list one row at a time and stacked back into a batch with np.stack, which cost more than the statevector simulation on cultivation d3 (about 60 ms of a 110 ms call of 262,144 shots). They now go into a two-batch array: a chunk adds at most one batch of survivors to fewer than one pending batch. Samples are bit-identical for a fixed seed. Quick benchmark, all threads: cultivation d3 2.61M -> 5.81M shots/s with trajectory sharing and 1.66M -> 2.56M with one lane per shot. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ChannelSampler.sample looped over every channel in Python and XOR-ed a dense uint8 row of all f-variables into each shot an error fired on; on surface d7r7 an error flips 2.7 of 337 f-variables, and the sampler took 91% of the time. A numba kernel now walks each channel's geometric skips and flips only the f-variables of the fired outcome. Blocks of 4096 shots run in parallel, each from its own SplitMix64 stream seeded with 64 bits from the sampler's Generator, so samples depend on the seed and the number of shots but not on the number of threads. The output layout is unchanged. The stream differs from NumPy's, so one test that pinned 4 of 10 events at p=0.3 now checks the rate over 100,000 shots. New tests cover a 2-bit channel's joint distribution across blocks, a certain error reaching every row of a partial last block, and reproducibility. Quick benchmark, all threads: surface d7r7 1.55M -> 22.4M shots/s, distillation (shared) 3.6M -> 9.0M, cultivation d5 (shared) 1.5M -> 2.7M. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With postselection, every chunk gathered the direct outputs into a zeroed array, copied its detector columns into the result and tested the postselected ones in three NumPy passes: 80 ms per 262,144 shots on cultivation d5 (47% of a call) and 466 ms on SymFT's d=7 circuit. One parallel numba pass now writes each direct detector into its result row and returns the discard mask (3.7 ms and 13 ms). Survivors are now buffered with one np.concatenate per chunk, as the review of the preallocated buffer suggested: as fast, without the two-batch invariant, and JAX never reads a buffer that is overwritten later. A new test checks that every survivor keeps its own direct detectors; the old tests could not catch misattributed rows, because their only direct detector was postselected. Samples are bit-identical for a fixed seed. Quick benchmark, all threads: distillation (shared) 9.0M -> 16.3M shots/s, cultivation d5 (shared) 2.7M -> 3.2M, SymFT d7 postselected (shared) 202k -> 241k. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- A channel whose identity probability is rounded below zero (e.g. PAULI_CHANNEL_1(0.33, 0.56, 0.11)) had a firing probability above 1, so log1p(-p) was NaN: undefined samples, a segfault on x86-64. The firing probability is now the clipped total of the non-identity outcomes, which is also exact for small probabilities. - Numba's default threading layer aborts the process when two Python threads run parallel code at once. tsim.utils.kernels.kernel compiles each kernel serially and in parallel, runs it in parallel only above a work threshold (waking the threads costs about 0.1 ms) and holds a lock around parallel calls. The parallel compilation gets its own name, since numba names cache files after the function and the two would share one entry. - Numba refuses to define a cached function where no cache directory is writable, which made `import tsim` fail on read-only installs; jit falls back to compiling without a cache. - SplitMix64 helpers move to tsim.utils.rng; geometric skips use a 64-bit uniform and return a float, inf when the uniform rounds to 1. - The direct-detector kernel takes named arguments built once per sampler. - Tests: blocks draw independent streams, a rounded identity probability, concurrent Python threads (in a subprocess), and a stronger survivor test with reference bits, flips, a shared f-bit, reordered outputs and a direct observable. A test whose circuit never reached the postselection path is removed; the reference test no longer counts the reference sample's row. - uv.lock adds only the numba lines (the earlier lock had been rewritten with a user-level index configuration). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The JAX shared executor kept static lane tables: kernels ran over every slot of a power-of-two table (26-40% live), the loop ran in segments with host checks, resized and resumed when a table filled, and compiled again for every new table size (about a second, randomly inside sample calls). On wide components that work dominated: 88% of a coherent d5r5 call. The executor is now a Python loop over the tape that keeps the lanes in growable NumPy arrays and calls small numba kernels: each action touches only the live lanes and the 2^extent amplitudes it can reach. extent comes from the plan (lower records one plus the highest active axis), and the planner now promotes onto the lowest free axis (a free SWAP of two dormant axes), so the extent tracks the active width: 1.8x less array work on coherent d5r5, 2.0x on SymFT's d=7 circuit. Rare outcomes are drawn as binomials by geometric skips with a 64-bit uniform, so the chance of a first hit stays exact below float32 resolution, as before. Grouping, splitting, pruning and reconstruction keep their semantics, and the exact distribution tests run unchanged against the new executor. The JAX executor stays for trajectory_sharing=False. shared_breakdown.py timed stages of the JAX executor and is removed. Quick benchmark, all threads, trajectory sharing: coherent d5r5 109k -> 436k shots/s, SymFT d7 postselected 241k -> 1.35M, cultivation d5 3.2M -> 8.6M, cultivation d3 5.1M -> 8.9M, distillation 16.3M -> 17.5M. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Once postselection had pruned every lane, the next active measurement took np.max of an empty array and raised; tiny batches, high rejection and zero shots hit it. The tape now stops when no lanes are left, and the norm deviation is a NaN-propagating max with an initial value, which the sampler now also warns about. - The lane table allocated at least 16 rows of 2^width amplitudes even for a few shots and copied full rows when growing. A lane holds at least one shot, so the table is capped at the shot count, and growing copies only the amplitudes the current extent reaches. - The JAX executor's per-lane shot count and choose_outcome's count argument existed only so that a future sharing executor could split lanes; both are gone. - Tests: a component whose extent drops and rises again (a generated d5r3 surface code with coherent over-rotations), sampled with NaN in every amplitude outside the tracked region, matches a run at full width bitwise; every kernel in parallel matches the serial run bitwise; shots whose lanes all die, and zero shots; every action stays within its extent. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- The extent test missed a clearing bug after the first rise: the generated circuit's low-extent stretches hold only rotations, so no lane moved while the extent was low. It now also runs with every fifth extent raised to the full width (an extent is only an upper bound), poisons the lane table with NaN after every split, prune and growth, and compares the norm deviations bitwise. - The parallel-vs-serial test's baseline ran the wide calls in parallel too; it now forces every kernel serial. - The all-dead test checks that some shot's lanes did all die. - A test checks that a NaN normalization deviation warns. - The lane table's shot cap no longer needs max(1, ...), and the module docstring is reflowed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
On the d=5 end-to-end cultivation circuit (463 qubits, 184,039 error bits) building a sampler took 101 s, 80 s of it in ChannelSampler: - np.unique(error_transform, axis=1) sorted all 184,039 columns of the 2,821-row error transform, comparing whole columns: 66 s. unique_columns finds the 15,148 distinct columns with a dict over the bit-packed columns and sorts only those; big-endian packing keeps the lexicographic order, so the result equals np.unique's. - absorb_subset_channels compared every pair of the 21,341 channels: 14 s. A channel of at most max_bits signatures has at most 15 strict subsets, so they are looked up in a dict, and absorbed in the original order. Both give the same noise model as before, bitwise, and the same samples for a fixed seed (checked on distillation, cultivation d3/d5, coherent d3r3 and the end-to-end circuit). Building the end-to-end sampler now takes 24 s, 17 s of it in zxparam's full_reduce. Tests compare unique_columns with np.unique and the indexed absorption with the pairwise one. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- unique_columns' docstring states the 1-D inverse (np.unique returned a 2-D one on numpy 2.0) and that nonzero entries count as 1, and drops benchmark figures that would go stale; `seen` is now `column_ids`. - The absorption test also draws channels without signatures and runs with max_bits 2 and 4; dropping the empty subset now fails it. - absorb_subset_channels' docstring and sort comment say what holds and why the order matters, and ChannelSampler's simplification steps are numbered 1 to 4. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Parallelism used to live inside each numba kernel, with serial Python, numpy and JAX glue between kernels, so sampling kept only 3-4 of the 14 cores of an M4 Pro busy. With trajectory_sharing=True the sampler now runs whole tasks on numba.get_num_threads() Python threads, and the kernels in them run serially: - utils/kernels.py: every numba function releases the GIL; run_tasks runs tasks on a thread pool and makes kernel() pick the serial compilation on its worker threads. Numba's cache ignores nogil, so jit tags the cached name with _nogil, as kernel() already did with _parallel. - ChannelSampler.block_seeds draws the seeds sample() drew itself, so the main thread draws them in order and hands them to the workers. - _sample_batches runs its batches on the threads. With postselection the shots are screened in chunks of at most 2^14 (noise, direct detectors, survivors), two chunks per thread at a time, until every thread has a full batch of survivors; the batches of survivors then run on the threads too. Seeds and keys are drawn in order and batches still start every batch_size survivors, so the samples do not depend on the thread count, and they equal the previous ones bit for bit when batch_size is at most 2^14 or a multiple of it (checked on cultivation d3, distillation, coherent d3r3, qv_q10 and the measurement sampler at 1, 3 and 14 threads). The JAX executors still sample one batch at a time. Against the previous commit, all threads, shots/s: workload (batch, shots per call) before after cultivation_d3 (65,536, 1M) 9.8M 19.6M coherent_d3_r3 (65,536, 1M) 6.4M 12.5M distillation (65,536, 1M) 19.9M 25.8M cultivation_d5 (65,536, 1M) 8.3M 10.7M coherent_d5_r5 (65,536, 1M) 393k 734k symft_msc_d7_postselected (1M) 2.71M 3.33M qv_q10 (16,384, 262k) 766k 1.60M Calls of a few batches gain less (1.0-1.6x at 262k shots), since one call has few survivor batches and their JAX and numpy glue holds the GIL. One thread is unchanged (within 7%). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With trajectory sharing every component runs in the numba executor on the host, but the sampler still moved each batch into JAX, gathered the direct outputs, concatenated and reordered the outputs with JAX ops, padded the last postselection survivor batch to batch_size rows (only JAX needs fixed shapes) and joined the batches on a JAX device. Now: - sample_program takes NumPy f-params and returns NumPy samples, put together by one numba kernel, _gather_outputs, that writes the direct outputs and the component outputs into their original columns. NumPy's column gathers were 4x slower than JAX's and ran under the GIL. - The sharing path passes NumPy f-params, samples the last survivor batch as it is, and joins its batches with np.concatenate. The JAX paths are unchanged. Samples without postselection are bit-identical to before. With postselection only the last survivor batch of a call differs: its padding rows were copies of row 0 that took part in the executor's grouping and random draws. A review compared the statistics with one lane per shot on 18 configurations, including calls whose every batch is partial (acceptance within |z| <= 1.6, matching outcome counts). Samples still do not depend on the thread count. Against the previous commit, all 14 threads, shots/s: workload (shots per call) before after symft_msc_d7_postselected (262k) 1.41M 2.07M symft_msc_d7_postselected (1M) 2.67M 3.64M distillation (262k) 20.1M 27.0M cultivation_d5 (262k) 9.1M 11.4M coherent_d3_r3 (1M) 12.9M 14.7M cultivation_d3 (1M) 21.6M 25.1M qv_q10 (262k) 1.51M 1.64M One thread: distillation +11%, cultivation_d5 +18%. Tests check that the host and JAX glue give equal outputs on every test circuit and on one with flipped direct detectors and observables, at 1, 300 and 70,000 rows, and that survivor batches hold exactly the surviving shots, without padding. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Calls of 2.6M shots on 10 threads still fell short of an in-process prototype that sampled one batch per thread: without trajectory sharing the JAX batches ran one at a time, and with it the main thread zeroed large outputs, every round started fresh threads, and Python made one GIL round trip per tape action. Now: - On a CPU backend the JAX statevector executor (one lane per shot) also samples a batch per numba thread, as many as fit in the memory budget, and the default batch size leaves room for them. The caller's jax.enable_x64, which is thread-local, reaches the worker threads, so precision and samples do not depend on the thread count. Stabilizer-rank programs still sample one batch at a time. - The noise and direct-detector kernels clear their own rows, so outputs come from np.empty: np.zeros of a large array is a serial memset when malloc reuses freed memory. - run_tasks keeps one thread pool (rebuilt in a forked child; no fork hook where os has none), waits for every task before raising, and runs nested calls inline. - The shared executor applies each stretch of consecutive rotations in one kernel call (qv_q10: 692 rotations in 11 calls), with the rotation coefficients computed once per layout. Samples are bit-identical to the previous commit with and without sharing, at 1 and 14 threads (sample hashes of cultivation d3, distillation, coherent d3r3, qv_q10 and the measurement sampler). shots/s, 10 threads with calls of 2.6M shots, 1 thread with calls of 262k; the prototype and the 10 single-thread processes ran the code before these changes: sharing on 1 thread 10 threads prototype 10 processes distillation 5.89M 41.2M 41.9M 41.8M coherent_d5_r5 147k 978k 979k 973k cultivation_d5 3.73M 20.3M 17.3M 22.8M qv_q10 737k 3.72M 1.82M 5.0M surface_d7_r7 12.6M 60.7M 73.3M 71.8M symft_msc_d7, postsel. 779k 3.18M - 2.85M cultivation_d3 7.63M 31.5M 26.2M 44.9M coherent_d3_r3 5.80M 19.6M - 37.4M sharing off distillation 4.25M 36.5M 26.2M 31.4M cultivation_d5 512k 4.04M 3.14M 3.0M cultivation_d3 1.12M 5.70M 5.34M 5.9M coherent_d3_r3 457k 2.30M 2.24M 2.28M surface_d7_r7 is bound by the OS freeing each call's 883 MB result (16 ms of a 45 ms call). On cultivation_d3 and coherent_d3_r3 a quarter to two fifths of each batch is executor bookkeeping in Python that holds the GIL. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The sampler's __init__ docstring had lost its `seed` entry (ruff D417). Two postselection tests rebuilt the drawn noise by appending each chunk's draw as it finished, but chunks are screened on several threads, so on CI's few-core runners the draws came back out of shot order. The tests now put the draws back in shot order by the order of their block seeds. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Without a batch size, a call is split into at least one batch per thread of at least MIN_THREAD_BATCH shots, and postselection splits the survivors left after the last chunk the same way (JAX pads them to a power of two). An explicit batch size keeps its samples independent of the thread count. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The shared executor's per-batch Python held the GIL for 20-30% of a batch on small circuits (the tape loop, lane split and prune in NumPy, group_on_host and the output assembly), which capped batch-per-thread sampling at 3-5 busy cores. `_execute` now samples a whole batch without the GIL: it groups the shots (a SplitMix64 Fisher-Yates shuffle, then stable sorts), runs the tape on lanes in arrays with a row per shot (only the amplitudes grow), and writes the outputs. Its parallel build splits a loop into chunks only once its work reaches PARALLEL_WORK, so a lone wide batch keeps the parallel kernels and small actions launch no threads. Shuffles and per-action seeds come from SplitMix64 instead of NumPy's Generator, so samples for a given seed change. One thread runs at the old speed (qv_q10 13.7M, SymFT d=7 postselected 1.39M against 1.29M shots/s). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The statevector strategy now runs entirely on the host: the numba executor samples both with trajectory sharing and with one lane per shot (share=False skips the grouping), so the JAX executor in execute.py and its tests are gone. StatevectorComponent is a frozen dataclass of NumPy arrays; batch seeds come from a NumPy Generator instead of JAX keys, and the memory budget is the host's. Every statevector batch gets a sharing layout, so postselected component detectors prune lanes in both modes. trajectory_sharing now defaults to True for active_statevector. The precision still follows jax_enable_x64 at compile time. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A call splits into a batch per thread only when every batch gets MIN_THREAD_BATCH shots; otherwise it stays one batch, which the parallel build of the executor spreads over the threads. On coherent d5r5, one batch of 16,384 shots samples 84.7k shots/s against 59k as four batches. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Components sampled overlapping random numbers: SplitMix64 counts up from its state, so seeds GOLDEN steps apart gave shifted copies of one stream. The executor now starts from the seed's mixed output; a test checks that two disconnected components are independent per shot and per batch. - With sharing the default, concurrent batches ignored the memory budget; _threads now applies the same fit with and without sharing. - The integration tests compare against strategy="cat5" instead of the default, which now shares too, and the plan-distribution tests run in complex128 as well and reject any outcome of probability 0. - Grouping keys are no longer computed without sharing, lane_bytes counts every per-shot array, and the docstrings no longer mention JAX, GPUs or the deleted execute module for the statevector path. - CHANGELOG entries for this PR. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
numba's parallel=True also turned array expressions such as np.zeros, np.full and slice fills into parallel loops, each a thread launch, so the parallel build of the shared executor launched threads at every measurement. Kernels now enable prange only (and get a new cache name). One SymFT d=7 postselected call of 16,384 shots samples 4.40M shots/s against 1.80M, cultivation d5 10.0M against 7.08M (65,536 shots). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- Postselection prunes in the statevector executor whenever a component detector is masked, not only when a direct detector is masked too. - Reference bits are XOR-ed into the outputs before a pruned shot's later outputs are cleared, so they stay False with reference samples as well. Outputs of f-bits only prune before the first action. - The sampler picks a JAX or a statevector backend once, instead of branching in every batching method, and the statevector path no longer creates JAX keys or queries JAX devices. CompiledProgram holds NumPy arrays. - Statevector amplitudes are complex128, independent of jax_enable_x64. - StatevectorComponent holds its arrays in one Tape; the executor takes the Tape and the SharingLayout as named tuples instead of positional tuples. - run_tasks uses one thread pool and a semaphore instead of a pool per thread count, and unique_columns accepts non-integer 0/1 matrices. - Smaller cleanups: fewer direct-output attributes, private sharing stats, one NaN-poisoning point for tests, docstrings and the changelog. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Summary
Stacked on #198 (base
trajectory-sharing), which is stacked on #203. This PR movesstrategy="active_statevector"off JAX and onto numba kernels on the CPU, and makes it use every core by default:share=Falseskips the grouping). The JAX executor (statevector/execute.py) is deleted.StatevectorComponentis a frozen dataclass of NumPy arrays, and batch seeds come from a NumPy Generator instead of JAX keys.active_statevectornow sample on the host; the stabilizer-rank strategies keep JAX and the GPU. Amplitudes are complex128, independent ofjax_enable_x64.active_statevector(trajectory_sharing=Nonemeans True there and False otherwise). In the benchmarks below it is never slower than one lane per shot.batch_size, a call is split into one batch per numba thread whenever each batch gets at least 4096 shots. Smaller calls stay one batch, which the parallel build of the executor spreads over the threads.batch_sizekeeps samples independent of the thread count.numba>=0.61is a runtime dependency.Changes in detail
statevector/shared.py)._execute: groups the shots (a SplitMix64 Fisher–Yates shuffle, then stable sorts on the key), runs the tape, and writes the outputs. Lanes live in arrays with one row per shot, which bounds them, so only the amplitudes grow.2^extentamplitudes it can reach. The extent comes from the plan, and the planner promotes onto the lowest free axis.PARALLEL_WORK. So a lone wide batch keeps parallel kernels, and small actions launch no threads.prangeloops. numba'sparallel=Truealso turnednp.zeros,np.fulland slice fills into thread launches at every measurement.ChannelSampler.sample): a numba kernel walks each channel's geometric skips and flips only the f-variables of the outcome that fired. Blocks of 4096 shots run in parallel, each on its own SplitMix64 stream, so samples depend on the seed and the shot count but not on the thread count. The random stream differs from NumPy's.np.concatenateper chunk.unique_columnsand an indexedabsorb_subset_channels. On the d=5 end-to-end cultivation circuit (463 qubits, 184,039 error bits), building the sampler takes 24 s instead of 101 s; 17 s of that is zxparam'sfull_reduce. The noise model is bitwise unchanged.tsim.utils.kernels:run_tasks).Performance (same machine, Apple M4 Pro, 14 cores)
Choose a thread setting and sharing mode; all threads with sharing (the default) is open.
All threads · trajectory sharing on (default)
All threads · trajectory sharing off
1 thread · trajectory sharing on (default)
1 thread · trajectory sharing off
How to read the tables:
clifft-benchmarks(benchmarks/clifft/comparison), run on 2026-10-19.sample(2_621_440, batch_size=65_536)). One thread meansNPROC=1 NUMBA_NUM_THREADS=1.trajectory-sharing(35db8b6).ActiveWidthSchedulePass. "Clifft k_max" is its peak active width, with the default passes' value in parentheses when it differs.With the default
batch_size=Noneon all threads, sampling one call of the harness's size, against the base of this stack's previous tip (ea4fc28):What the numbers show:
ActiveWidthSchedulePassreduces to width 0, and with all threads on distillation, cultivation d3 and coherent d3r3. Clifft'ssample_survivorsreturns counts only, while Tsim's harness builds the full detector arrays, which costs Tsim on the cheapest circuits.Changes from the stack review (7e13528)
CompiledProgramholds NumPy arrays.Tapeand theSharingLayoutare passed as named tuples instead of positional tuples, andStatevectorComponentholds its arrays in oneTape.run_tasksuses one thread pool and a semaphore instead of a pool per thread count.unique_columnsaccepts non-integer 0/1 matrices.How this was verified
strategy="cat5"on every test circuit and on random noisy circuits.batch_sizedon't depend on it.unique_columnsagainstnp.unique, bitwise-identical noise models).Test plan
pytest: 1715 passedpyright src test: 0 errorsKnown limitations, left for follow-ups:
🤖 Generated with Claude Code