Skip to content

Run strategy="active_statevector" in numba kernels, without JAX - #199

Merged
rafaelha merged 21 commits into
trajectory-sharingfrom
tsim-perf
Oct 20, 2026
Merged

rafaelha merged 21 commits into
trajectory-sharingfrom
tsim-perf

Conversation

@rafaelha

@rafaelha rafaelha commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Stacked on #198 (base trajectory-sharing), which is stacked on #203. This PR moves strategy="active_statevector" off JAX and onto numba kernels on the CPU, and makes it use every core by default:

  • No JAX in the statevector path.
    • One numba function samples a whole batch with the GIL released, both with trajectory sharing and with one lane per shot (share=False skips the grouping). The JAX executor (statevector/execute.py) is deleted.
    • StatevectorComponent is a frozen dataclass of NumPy arrays, and batch seeds come from a NumPy Generator instead of JAX keys.
    • Behavior change: GPU users of active_statevector now sample on the host; the stabilizer-rank strategies keep JAX and the GPU. Amplitudes are complex128, independent of jax_enable_x64.
  • Trajectory sharing is the default for active_statevector (trajectory_sharing=None means True there and False otherwise). In the benchmarks below it is never slower than one lane per shot.
  • Every thread has work.
    • Without a 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.
    • With postselection, the survivors left after screening are split the same way.
    • An explicit batch_size keeps samples independent of the thread count.
  • Noise sampling and setup (from the earlier commits of this PR) run in numba kernels as well; numba>=0.61 is a runtime dependency.

Changes in detail

  • Shared executor (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.
    • Small circuits: before, 20–30% of a batch was GIL-held Python and NumPy, which capped batch-per-thread sampling at 3–5 busy cores.
    • Wide components: each action touches only the live lanes and the 2^extent amplitudes it can reach. The extent comes from the plan, and the planner promotes onto the lowest free axis.
    • Parallel build: a loop is split into chunks across threads only once its work reaches PARALLEL_WORK. So a lone wide batch keeps parallel kernels, and small actions launch no threads.
    • Kernels parallelize only explicit prange loops. numba's parallel=True also turned np.zeros, np.full and slice fills into thread launches at every measurement.
    • Randomness: rare outcomes are drawn by geometric skips with a 64-bit uniform, so they stay exact below float32 resolution. Every stream starts from its seed's SplitMix64 output, so different components never share random numbers. The samples for a given seed change.
    • Postselection: postselected component detectors now prune lanes with sharing off as well.
  • Noise sampling (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.
  • Direct detectors and postselection: one parallel numba pass writes each direct detector into its result row and returns the discard mask. Survivors are buffered with one np.concatenate per chunk.
  • Sampler setup on large circuits: unique_columns and an indexed absorb_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's full_reduce. The noise model is bitwise unchanged.
  • tsim.utils.kernels:
    • Serial or parallel: compiles every kernel serially and in parallel, and runs the parallel version only above a work threshold.
    • Locking: holds a lock around parallel calls, because the workqueue layer aborts on concurrent ones.
    • Worker threads: runs whole batches on a cached thread pool (run_tasks).
    • Read-only installs: compiles without a cache there.

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)
circuit Tsim width Clifft k_max base this PR Clifft this PR / base this PR / Clifft
surface d7r7† 0 0 1.59M 30.7M 37.8M 19× 0.81×
coherent d3r1† 0 0 (4) 93.6M 103M 682M 1.1× 0.15×
coherent d5r1† 0 0 (12) 66.3M 90.5M 447M 1.4× 0.2×
distillation 3 3 (5) 3.21M 32.9M 37.1M 10× 0.89×
cultivation d3 4 4 2.24M 30.5M 46.3M 14× 0.66×
coherent d3r3 4 4 (5) 1.73M 21.3M 46.7M 12× 0.46×
cultivation d5 6 10 1.25M 16M 1.65M 13× 9.7×
coherent d5r5 11 13 18.4k 155k 88.5k 8.5× 1.8×
quantum volume, 10 qubits 10 10 837k 93M 174k 111× 535×
quantum volume, 20 qubits 20 20 8.1 154 17 19× 9.1×
SymFT d=7, no postselection 8 19 11.2k 475k 1.38k 42× 344×
SymFT d=7, postselected 8 19 423k 6.99M 38.7k 17× 181×
All threads · trajectory sharing off
circuit Tsim width Clifft k_max base this PR Clifft this PR / base this PR / Clifft
surface d7r7 0 0 1.59M 30.7M 37.8M 19× 0.81×
coherent d3r1 0 0 (4) 93.6M 103M 682M 1.1× 0.15×
coherent d5r1 0 0 (12) 66.3M 90.5M 447M 1.4× 0.2×
distillation 3 3 (5) 3.13M 26.3M 37.1M 8.4× 0.71×
cultivation d3 4 4 2.06M 10.1M 46.3M 4.9× 0.22×
coherent d3r3 4 4 (5) 1.27M 7.11M 46.7M 5.6× 0.15×
cultivation d5 6 10 733k 4.48M 1.65M 6.1× 2.7×
coherent d5r5 11 13 4.47k 151k 88.5k 34× 1.7×
quantum volume, 10 qubits 10 10 5.35k 13.7k 174k 2.6× 0.079×
quantum volume, 20 qubits 20 20 0.99 1.8 17 1.9× 0.11×
SymFT d=7, no postselection 8 19 53.8k 450k 1.38k 8.4× 326×
SymFT d=7, postselected 8 19 130k 2.97M 38.7k 23× 77×
1 thread · trajectory sharing on (default)
circuit Tsim width Clifft k_max base this PR Clifft this PR / base this PR / Clifft
surface d7r7† 0 0 1.54M 11.2M 6.18M 7.3× 1.8×
coherent d3r1† 0 0 (4) 90.4M 108M 682M 1.2× 0.16×
coherent d5r1† 0 0 (12) 63.9M 79.6M 275M 1.2× 0.29×
distillation 3 3 (5) 3.16M 5.98M 3.45M 1.9× 1.7×
cultivation d3 4 4 2.01M 8.13M 4.98M 4.1× 1.6×
coherent d3r3 4 4 (5) 1.65M 5.91M 5.11M 3.6× 1.2×
cultivation d5 6 10 1.09M 3.74M 172k 3.4× 22×
coherent d5r5 11 13 6.29k 37.4k 9.55k 5.9× 3.9×
quantum volume, 10 qubits 10 10 853k 16.2M 25.7k 19× 629×
quantum volume, 20 qubits 20 20 2.5 39 4.4 15× 8.8×
SymFT d=7, no postselection 8 19 8.04k 58.1k 200 7.2× 291×
SymFT d=7, postselected 8 19 398k 1.7M 7.18k 4.3× 237×
1 thread · trajectory sharing off
circuit Tsim width Clifft k_max base this PR Clifft this PR / base this PR / Clifft
surface d7r7 0 0 1.54M 11.2M 6.18M 7.3× 1.8×
coherent d3r1 0 0 (4) 90.4M 108M 682M 1.2× 0.16×
coherent d5r1 0 0 (12) 63.9M 79.6M 275M 1.2× 0.29×
distillation 3 3 (5) 2.8M 5.03M 3.45M 1.8× 1.5×
cultivation d3 4 4 1.3M 1.63M 4.98M 1.3× 0.33×
coherent d3r3 4 4 (5) 669k 1.13M 5.11M 1.7× 0.22×
cultivation d5 6 10 432k 788k 172k 1.8× 4.6×
coherent d5r5 11 13 1.36k 21.7k 9.55k 16× 2.3×
quantum volume, 10 qubits 10 10 1.66k 1.45k 25.7k 0.87× 0.056×
quantum volume, 20 qubits 20 20 0.33 0.26 4.4 0.78× 0.059×
SymFT d=7, no postselection 8 19 23.8k 57.2k 200 2.4× 287×
SymFT d=7, postselected 8 19 73.1k 1.03M 7.18k 14× 144×

How to read the tables:

  • Units: attempted shots/s, counted before postselection.
  • Source: the harness on clifft-benchmarks (benchmarks/clifft/comparison), run on 2026-10-19.
  • Configurations: each Tsim column shows its best configuration from the harness calls. With all threads, that includes calls of 40 batches (for example sample(2_621_440, batch_size=65_536)). One thread means NPROC=1 NUMBA_NUM_THREADS=1.
  • "base": the same harness on trajectory-sharing (35db8b6).
  • Clifft: 0.11.0 with every pass, including the opt-in ActiveWidthSchedulePass. "Clifft k_max" is its peak active width, with the default passes' value in parentheses when it differs.
  • †: no statevector component, so sharing doesn't apply.

With the default batch_size=None on all threads, sampling one call of the harness's size, against the base of this stack's previous tip (ea4fc28):

circuit (shots per call) sharing on: before → now sharing off: before → now
distillation (262,144) 30.8M → 33.7M 17.5M → 20.5M
cultivation d3 (262,144) 22.3M → 44.0M 9.47M → 11.0M
coherent d3r3 (262,144) 17.4M → 31.9M 4.52M → 8.63M
cultivation d5 (65,536) 7.74M → 9.24M 450k → 1.87M
coherent d5r5 (16,384) 76.5k → 87.8k 6.5k → 38.8k
quantum volume, 10 qubits (65,536) 24.7M → 37.3M 9.4k → 13.4k
SymFT d=7, postselected (16,384) 2.76M → 4.18M 76k → 1.83M

What the numbers show:

  • Against base: this PR is faster on every circuit, except one lane per shot on the two quantum-volume circuits on one thread (0.87× and 0.78×). There the deleted JAX executor was faster; with sharing, the default, the same circuits are 15–19× faster than base.
  • Against Clifft with all passes:
    • Tsim wins on every circuit with a non-trivial statevector component from cultivation d5 on: 1.8× to 535× on all threads with sharing.
    • On one thread with sharing, Tsim is also ahead on distillation, cultivation d3 and coherent d3r3.
    • Clifft is faster on the coherent r1 circuits, which ActiveWidthSchedulePass reduces to width 0, and with all threads on distillation, cultivation d3 and coherent d3r3. Clifft's sample_survivors returns counts only, while Tsim's harness builds the full detector arrays, which costs Tsim on the cheapest circuits.
  • d=5 end-to-end cultivation, postselected on the 107 cultivation-stage detectors, measured at af98f77, before the executor and batching commits: this PR 620k (all threads) and 206k (1 thread) shots/s; Clifft 0.11 542k and 70k.

Changes from the stack review (7e13528)

  • Pruning: postselection prunes in the executor whenever a component detector is masked, also when no direct detector is.
  • Reference samples: reference bits are XOR-ed into the outputs before a pruned shot's later outputs are cleared, so those stay False. Outputs of f-bits only prune before the first action.
  • Backends: the sampler picks a JAX or a statevector backend once, instead of branching in every batching method. The statevector path creates no JAX keys and queries no JAX devices, and CompiledProgram holds NumPy arrays.
  • Precision: statevector amplitudes are complex128.
  • Executor arguments: the Tape and the SharingLayout are passed as named tuples instead of positional tuples, and StatevectorComponent holds its arrays in one Tape.
  • Threads: run_tasks uses one thread pool and a semaphore instead of a pool per thread count.
  • Noise setup: unique_columns accepts non-integer 0/1 matrices.

How this was verified

  • Executor tests:
    • Exact distribution: sharing on and off; an outcome of probability 0 is never drawn.
    • Pruning: exact pruning distributions, and outputs set after a lane's pruning action read False.
    • Rare outcomes: binomial counts below float32 resolution.
    • NaN poisoning: the amplitudes the executor must not read hold NaN while a d5r3 surface code whose extent drops and rises is sampled. It matches a full-width run bitwise.
    • Parallel and serial builds give the same samples.
  • Sampler tests:
    • Baselines: sharing against strategy="cat5" on every test circuit and on random noisy circuits.
    • Independence: two disconnected components draw independent numbers per shot and per batch. A review found the first version gave them shifted copies of one SplitMix64 stream.
    • Thread count: samples with an explicit batch_size don't depend on it.
    • Batching: default batches give each thread one; postselection splits the survivors.
    • Memory: the memory fit applies with and without sharing.
  • Noise sampling, direct detectors and setup: as in the earlier commits (joint distributions across blocks, reproducibility, concurrent threads, unique_columns against np.unique, bitwise-identical noise models).
  • Harness check: acceptance and observable rates agree with Clifft on all 12 workloads (|z| ≤ 2.5).
  • Review: a review agent checked every step. Its fixes are the "Fix …", "Strengthen …" and "Polish …" commits.

Test plan

  • pytest: 1715 passed
  • pyright src test: 0 errors

Known limitations, left for follow-ups:

  • Sharing on all threads keeps 5–7 of 14 cores busy on small circuits. The rest is serial work around the batches: noise sampling, screening and concatenation.
  • Memory-limited concurrency: with complex128, the per-shot memory estimate of a statevector batch doubles, so on a machine with little free memory fewer batches run at once. With about 17 GiB free, quantum volume q10 runs 6 instead of 13 batches at once (114M instead of 154–163M shots/s), and postselected coherent d5r5 runs 1 instead of 2 (88k instead of 96–99k shots/s). The kernels themselves run as fast as with complex64. The tables above were measured with complex64.

🤖 Generated with Claude Code

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Oct 14, 2026 •

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-10-14T01:32:43.833920Z e394e84 PR opened
ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@github-actions

github-actions Bot commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://queracomputing.github.io/tsim/pr-preview/pr-199/

Built to branch gh-pages at 2026-10-20 15:01 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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]

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

@github-actions

github-actions Bot commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor

☂️ Code Coverage

current status: ✅

Overall Coverage

Statements Covered Coverage Threshold Status
4105 3604 88% 0% 🟢

New Files

File Coverage Status
src/tsim/utils/kernels.py 93% 🟢
src/tsim/utils/rng.py 47% 🟢
TOTAL 70% 🟢

Modified Files

File Coverage Status
src/tsim/circuit.py 92% 🟢
src/tsim/compile/pipeline.py 100% 🟢
src/tsim/core/types.py 100% 🟢
src/tsim/noise/channels.py 93% 🟢
src/tsim/sampler.py 94% 🟢
src/tsim/statevector/_init_.py 100% 🟢
src/tsim/statevector/plan.py 99% 🟢
src/tsim/statevector/program.py 100% 🟢
src/tsim/statevector/shared.py 26% 🟢
TOTAL 89% 🟢

updated for commit: 7e13528 by action🐍

@rafaelha
rafaelha added this pull request to stack #200 October 14, 2026 02:09
@rafaelha
rafaelha removed this pull request from stack #200 October 19, 2026 19:24
@rafaelha
rafaelha added this pull request to stack #204 October 19, 2026 19:24
@rafaelha rafaelha changed the title Run noise sampling and trajectory sharing in numba kernels Run strategy="active_statevector" in numba kernels, without JAX Oct 19, 2026
rafaelha and others added 16 commits October 20, 2026 10:03
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>
rafaelha and others added 5 commits October 20, 2026 10:03
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>
@rafaelha
rafaelha merged commit 2b17240 into main Oct 20, 2026
10 of 11 checks passed
@rafaelha
rafaelha deleted the tsim-perf branch October 20, 2026 15:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant