Run noise sampling and trajectory sharing in numba kernels by rafaelha · Pull Request #199 · QuEraComputing/tsim

rafaelha · 2026-10-14T01:25:44Z

Summary

Stacked on #198 (base trajectory-sharing), which is stacked on #196. On that branch, two costs dominated:

Small circuits: with trajectory sharing on, the classical per-shot work dominated the call.
Wide components: the JAX shared executor spent most of each call (88% on coherent d5r5) on static-table overhead.

This PR moves both onto numba kernels on the CPU. It also makes building a sampler for large circuits 4× faster. numba is now a runtime dependency (numba>=0.61).

Noise sampling (ChannelSampler.sample):
- Before: a Python loop over channels XOR-ed a dense row into every shot an error fired on.
- Now: a numba kernel walks each channel's geometric skips and flips only the f-variables of the outcome that fired.
- Reproducibility: blocks of 4096 shots run in parallel, each on its own SplitMix64 stream seeded from the sampler's Generator. Samples depend on the seed and the shot count, but not on the thread count.
- Behavior change: the random stream differs from NumPy's, so the samples for a given seed change.
Direct detectors and postselection:
- One parallel numba pass writes each direct detector into its result row and returns the discard mask. On cultivation d5 this takes 3.7 ms instead of 80 ms per 262,144 shots.
- Survivors are buffered with one np.concatenate per chunk instead of row by row in a Python list.
Trajectory-sharing executor:
- Before: a JAX executor over static lane tables.
- Now: a Python loop over the tape. It 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: extent comes from the plan. The planner now promotes onto the lowest free axis, so the extent tracks the active width. This does 1.8× less array work on coherent d5r5 and 2.0× less on SymFT's d=7.
- No recompiles: there is no recompilation per table size.
- Rare outcomes: drawn by geometric skips with a 64-bit uniform, so they stay exact below float32 resolution.
- What stays: the JAX executor is still used for trajectory_sharing=False. Its per-lane count and the count argument of choose_outcome existed only for sharing, and both are removed.
Sampler setup on large circuits:
- unique_columns builds a dict over the bit-packed columns. Its result equals np.unique(axis=1).
- An indexed absorb_subset_channels looks up the at most 15 strict subsets of each channel instead of comparing all pairs.
- Effect: 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: this module handles three numba pitfalls:
- Serial or parallel: it compiles every kernel both ways and runs the parallel version only above a work threshold, because a parallel launch costs about 0.13 ms on macOS.
- Locking: it holds a lock around parallel calls, because numba's workqueue layer aborts the process on concurrent parallel calls from Python threads.
- Read-only installs: it compiles without a cache there.
tsim.utils.rng holds the SplitMix64 helpers.
Fixed along the way: a channel whose identity probability rounds below zero, such as PAULI_CHANNEL_1(0.33, 0.56, 0.11), had a firing probability above 1, which made the geometric skips NaN. The firing probability is now the clipped total of the non-identity outcomes.

Performance (same machine, Apple M4 Pro)

circuit	base, 1 thread	this PR, 1 thread	Clifft, 1 thread	base, all threads	this PR, all threads	Clifft, all threads
surface d7r7	1.56M	10.2M	6.25M	1.57M	24.6M	40.4M
distillation	3.27M	5.41M	2.89M	3.28M	16.5M	25.7M
cultivation d3	2.37M	7.05M	5.00M	2.50M	9.52M	39.3M
coherent d3r3	1.76M	5.57M	2.03M	1.66M	6.93M	17.0M
cultivation d5	1.19M	3.15M	150k	1.25M	7.98M	1.37M
coherent d5r5	42.5k	146k	3.8k	103k	425k	34.4k
SymFT d=7, postselected	98.5k	576k	1.7k	184k	1.50M	7.3k
SymFT d=7, no postselection	909	2.1k	86	2.5k	13.2k	381
quantum volume, 10 qubits	392k	700k	30.5k	523k	851k	186k

How to read the table:

Units: attempted shots/s, counted before postselection.
Source: scripts/clifft_comparison/run_all.sh, with Clifft 0.11.0.
Threads: "1 thread" means NPROC=1 NUMBA_NUM_THREADS=1.
Configurations: each Tsim column shows its best configuration, one lane per shot or shared.
"base": the same harness run on trajectory-sharing.
When it was measured: at af98f77, which is this branch before the rebase onto Give graph building one owner: GraphBuilder #197 and before the two noise-setup commits. Those commits leave the noise model and the samples bitwise unchanged.

What the numbers show:

One thread: Tsim now beats Clifft on every circuit in the table. Clifft is still faster on 20-qubit quantum volume (5 against 2 shots/s).
All threads:
- Clifft is still 1.6× to 4.1× faster on surface d7r7, distillation, cultivation d3 and coherent d3r3.
- Tsim is faster on the rest: 5.8× on cultivation d5, 12× on coherent d5r5, 35× and 205× on SymFT's d=7 without and with postselection, and 4.6× on 10-qubit quantum volume.
d=5 end-to-end cultivation, postselected on the 107 cultivation-stage detectors:

all threads 1 thread

before (on Add strategy="statevector" for sampling on a dense state vector #196) not measured 12.9k

this PR 620k 206k

Clifft 0.11 542k 70k

How this was verified

Shared executor: the exact distribution tests of the base PR run unchanged against the numba executor. New tests:
- Extent drop and rise: a generated d5r3 surface code with coherent over-rotations, whose extent drops and rises again, is sampled with NaN in every amplitude outside the tracked region. It matches a full-width run bitwise.
- Parallel vs serial: every kernel run in parallel matches the serial run bitwise.
- Edge cases: shots whose lanes all die, zero shots, and every action staying within its extent.
Noise sampling:
- a 2-bit channel's joint distribution across blocks;
- a certain error reaching every row of a partial last block;
- reproducibility;
- concurrent Python threads, run in a subprocess;
- a rounded identity probability.
Direct detectors: every survivor keeps its own direct detectors. This was checked with reference bits, flips, a shared f-bit, reordered outputs and a direct observable.
Setup:
- unique_columns equals np.unique.
- The indexed absorption equals the pairwise one.
- The noise model is bitwise identical, and samples for a fixed seed are the same, on distillation, cultivation d3 and d5, coherent d3r3 and the end-to-end circuit.
Harness check: acceptance and observable rates agree with Clifft on all 12 workloads. The largest deviation is |z| = 2.3 on coherent d5r5 acceptance, as on the base branch.
Review: a review agent checked every step. Its fixes are the "Fix …", "Strengthen …" and "Polish …" commits.

Test plan

pytest: 1641 passed
pyright src test: 0 errors
pre-commit hooks pass on the changed files

Known limitations, left for follow-ups:

Multi-core scaling: all threads run only 1.2× to 6× faster than one thread.
- On macOS, numba has only the workqueue threading layer: TBB has no macOS wheels, and OpenMP needs libomp.
- A parallel launch costs 65–120 µs, and the shared executor launches one or two kernels per tape action.
- With OpenMP the executor gains only 1.1× to 1.6×, because the serial NumPy lane split and prune take about half of the d=7 executor's time.
Trajectory sharing runs on the CPU only. On a GPU, use trajectory_sharing=False.
Setup on the end-to-end circuit: the remaining 17 s is zxparam's full_reduce. 187k of its 225k vertices are noise spiders.

🤖 Generated with Claude Code

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>

The one-thread Tsim run set only NPROC=1, which pins XLA; numba (noise sampling, trajectory sharing) now also gets NUMBA_NUM_THREADS=1. Shared samplers no longer recompile for new lane tables, so every configuration gets a single warmup call. 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>

chatgpt-codex-connector · 2026-10-14T01:25:56Z

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 · 2026-10-14T01:26:59Z

PR Preview Action v1.8.1
🚀 View preview at https://queracomputing.github.io/tsim/pr-preview/pr-199/
Built to branch `gh-pages` at 2026-10-14 01:26 UTC. Preview will be ready when the GitHub Pages deployment is complete.

chatgpt-codex-connector

💡 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".

chatgpt-codex-connector · 2026-10-14T01:32:41Z

+    _, 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]


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 · 2026-10-14T01:32:43Z

☂️ Code Coverage

current status: ✅

Overall Coverage

Statements	Covered	Coverage	Threshold	Status
3864	3574	92%	0%	🟢

New Files

File	Coverage	Status
src/tsim/utils/kernels.py	91%	🟢
src/tsim/utils/rng.py	47%	🟢
TOTAL	69%	🟢

Modified Files

File	Coverage	Status
src/tsim/noise/channels.py	93%	🟢
src/tsim/sampler.py	92%	🟢
src/tsim/statevector/execute.py	100%	🟢
src/tsim/statevector/plan.py	100%	🟢
src/tsim/statevector/program.py	100%	🟢
src/tsim/statevector/shared.py	67%	🟢
TOTAL	92%	🟢

updated for commit: e394e84 by action🐍

rafaelha and others added 10 commits October 13, 2026 11:40

chatgpt-codex-connector Bot reviewed Oct 14, 2026

View reviewed changes

rafaelha added this pull request to stack #200 October 14, 2026 02:09

Navigation Menu

Repository navigation

Run noise sampling and trajectory sharing in numba kernels - #199

Run noise sampling and trajectory sharing in numba kernels#199
rafaelha wants to merge 10 commits into
trajectory-sharingfrom
tsim-perf

rafaelha commented Oct 14, 2026

Uh oh!

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

edited

Loading

Uh oh!

github-actions Bot commented Oct 14, 2026

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

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Oct 14, 2026

Uh oh!

github-actions Bot commented Oct 14, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

	all threads	1 thread
before (on Add strategy="statevector" for sampling on a dense state vector #196)	not measured	12.9k
this PR	620k	206k
Clifft 0.11	542k	70k

Conversation

rafaelha commented Oct 14, 2026

Summary

Performance (same machine, Apple M4 Pro)

How this was verified

Test plan

Uh oh!

chatgpt-codex-connector Bot commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codex Review Summary

Uh oh!

github-actions Bot commented Oct 14, 2026

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

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Oct 14, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 14, 2026

☂️ Code Coverage

Overall Coverage

New Files

Modified Files

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

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

edited

Loading

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