Repository navigation
perf(metal): select the requested output columns on the GPU - #294
Conversation
Detector samples that return a subset or reordering of the outputs
(separate_observables, prepend_observables, or detectors only) split the
rows on the host: bit-packed columns not leading the rows were unpacked
to booleans, indexed and repacked (rows.select), and boolean columns
were copied out of every output unpacked. On distance-3 cultivation,
postselected on every detector, that took 0.48 s of a 0.53 s call of
2^26 shots on an M4 Pro.
A new select_rows kernel writes each requested column list into its own
array, bit-packed or as booleans, one thread per output byte, in one
command buffer. All outputs in order stay the rows (or unpack_rows), and
leading whole bytes a view of them; calls past 32-bit indices keep the
host path. Boolean arrays stay C-contiguous.
M shots/s on an M4 Pro, main -> this commit (postselected as recorded;
the GPU does the work, so the numba thread count does not matter):
cultivation d3 (21 outputs) bit-packed booleans
separate observables 147 -> 1316 97 -> 1104
prepended observables 55 -> 1508 89 -> 1249
detectors only 1312 -> 1468 101 -> 1241
appended observables 1835 -> 1863 1424 -> 1391
end2end d5 (2745 outputs)
separate observables 59 -> 291 3 -> 156
prepended observables 1 -> 221 3 -> 171
detectors only 305 -> 306 3 -> 170
appended observables 308 -> 307 214 -> 213
Without postselection the same layouts gain alike (cultivation 21-102
-> 888-1187, end2end 0-7 -> 14-65), where end2end's rows of 344 bytes
bound every layout near 68 M/s.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Codex Review Summary: ✅ CompletedThis 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. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a913858530
ℹ️ 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".
| if total * width >= 1 << 32: | ||
| selected.append(_select(packed, chosen, num_outputs, not bools)) |
There was a problem hiding this comment.
When a boolean selection is large enough that total * width >= 2**32, this fallback calls gpu.rows.select with bit_packed=False; that helper performs advanced column indexing (bits[:, columns]), which produces an F-contiguous rather than C-contiguous NumPy array. Thus sufficiently large detector-only, prepended, or separate-observable calls violate the C-contiguous behavior preserved by the normal GPU path and can break consumers that require C-order buffers. Make the boolean fallback contiguous as the removed _columns helper did.
Useful? React with 👍 / 👎.
☂️ Code Coverage
Overall Coverage
New FilesNo new covered files... Modified Files
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Review of #294: a copy kernel's thread steps by the dispatch's threads (up to 2^20), so with an output just under the 2^32-byte guard its index could pass 2^32, wrap to where it started and never end; the GPU would stall until Metal's watchdog failed the command buffer. select_rows (bit-packed column subsets, now on the GPU) and unpack_rows (already on main) fall back to the host from MAX_COPY_ITEMS = 2^32 - 2^20 bytes on. That fallback's booleans are now C-contiguous, as the GPU's are, and a test forces it. Also: the CHANGELOG's ranges match the measurements, a test comment names the shapes whose observables end inside a byte, the host helper is imported as _host_select so that it does not read as the method _select, and sample_accepted's Returns names the list it gives with columns. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Problem
On an Apple GPU, detector samples that return a subset or reordering of the outputs split each call's rows on the host:
separate_observables=True, every column ofprepend_observables=True) went throughgpu.rows.select, which unpacks every row to booleans, indexes the columns and repacks them._columns).The GPU idled meanwhile. On distance-3 cultivation postselected on every detector,
sample(2**26, separate_observables=True, bit_packed=True)spent 0.48 s of its 0.53 s inrows.select, and the GPU work took 0.04 s. The most Stim-like call,detectors, observables = sampler.sample(..., separate_observables=True), was 10 to 15 times slower than appended observables. On end-to-end cultivation d=5 (2745 outputs), several layouts sampled 0 to 3 M shots/s.Fix
A new Metal kernel,
select_rows, writes each requested column list into its own array, bit-packed (zero bits past the columns) or as one boolean per output. It runs one thread per output byte, and the column lists of a call share one command buffer.unpack_rows(booleans), and leading whole bytes stay a view of the rows._columnsis gone, andMetalShotSampler.sample_acceptedtakes the column lists (columns=) and returns one array per list.Numbers
Apple M4 Pro, 20-core GPU, Metal. The work is on the GPU, so the numba thread count does not change these.
main(7ec136e6) → this branch, in M shots/s, postselected as tsim-bench records:Without postselection the same layouts gain alike: cultivation goes from 21–102 to 888–1187, and end2end from 0–7 to 14–65. Every layout of end2end is then bound near 68 M shots/s by its 344-byte rows.
These grids come from a timing script, because
tsim-dev benchonly times appended observables. That bench is unchanged:tsim-dev bench --device gpu --ref origin/main, M4 Pro, 1 numba thread, quick suite:Tests
test_metal_selects_every_layout_of_columnschecks separate, prepended and detectors-only output, bit-packed and as booleans, with and without postselection, against appended booleans (the unchangedunpack_rowspath). Its five shapes have observables that start inside a byte, end inside one, span two bytes, or align with bytes. It also checks that booleans are C-contiguous.test_metal_samples_no_shotsnow also checksseparate_observables=Truewith zero shots, which an early version of this change returned as one array instead of two.Verification
uv run tsim-dev checkpasses, apart fromtest_dev_skills_have_valid_frontmatter. That one fails only because my checkout has untracked local skills in.claude/skills, so CI will not see it.gpu/rows.pyis unchanged.🤖 Generated with Claude Code