Repository navigation
simd: branch-free mask_gather_u32 + gated mask_gather_u32_under - #342
Conversation
mask_gather_u32 tested `idx < src_rows && bit == 1` per row. Whether that compiles to a branch depends on the call site. Called through lance-graph's mask-risc Gather on a 1M-row lane with an unpredictable source, it ran 5.3 ms. The branch-free form runs 1.5 ms, level with a hand-written branch-free loop (lance-graph D-GATED-GATHER-0, re-run against this commit). Inlined into a plain loop, the compiler had already removed the branch, and the two forms time the same. The new form shifts each row's bit in unconditionally; the out-of-range test is a select on the address plus an AND on the bit (gather_bit). Same loads, same contract. mask_gather_u32_under(src, src_rows, index, under, out) returns gather & under, but reads index only at set gate bits. It skips zero gate words and walks set bits with trailing_zeros, so the cost tracks popcount(under). The same probe measured this schedule fastest, or within noise of fastest, at every gate density from 0.01% to 100%. Both gathers return early on src_rows == 0, because src may then be empty and gather_bit's fallback read of src[0] would not exist. Gate bits past index.len() are cleared before they can address index. Tests: parity against a naive reference and against gather & gate, across word tails, with dirty output buffers and dirty gate tails. The simd-masking-parity crate gains check 0xD01 for the gated form. Disable runs, each red and then restored: dropping the gate-tail clamp panics on an out-of-bounds index read; dropping the gated form's src_rows == 0 return panics on an empty src. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HdJxpkHATkdNL2veorKao2
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (3)
Included review availability: This review used your included allowance. 0 included reviews remain after this review. Your included PR review attempts over the past 7 days set your current allowance at 2 reviews per hour. 📝 WalkthroughWalkthrough
ChangesSIMD Masked Gather
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Feature Suggested reviewers: Merge Risk: ⚪ Minimal · up to No confirmed issue prevents merging after normal checks. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
Warning Billing warning: we have not been able to collect payment for this subscription for more than 72 hours. Please update the payment method or pay any pending invoices in Billing to avoid service interruption. A rabbit checks each gathered bit, Comment |
Bugbot couldn't run - usage limit reachedBugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit. A user or team admin can review and increase usage limits in the Cursor dashboard. (requestId: 2501f148-2933-436d-b351-043194abb7cc) |
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. |
Bugbot couldn't run - usage limit reachedBugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit. A user or team admin can review and increase usage limits in the Cursor dashboard. (requestId: 09d42580-55b8-4e52-a227-bd7a940ae996) |
2 participants
What
mask_gather_u32is branch-free per row. Each row's bit is shifted in unconditionally. The out-of-range test is a select on the address plus an AND on the bit (gather_bit). The loads and the contract are unchanged.mask_gather_u32_under(src, src_rows, index, under, out). It returnsgather & under, but readsindexonly at set gate bits. It skips zero gate words and walks set bits withtrailing_zeros, so the cost followspopcount(under), notindex.len(). It is exported through thendarray::simdfacade.Why
lance-graph
D-GATED-GATHER-0found two costs in the semijoingate ∧ gather(fk, foreign):if idx < rows && bit == 1compiles to a branch depends on the call site.Measured
I re-ran lance-graph's
gated_gather_probeagainst this commit: 1M rows, a 1024-row foreign table with half its rows set, release build, debug info off, median of 15. Every route's count is asserted equal to an oracle.All times are µs.
Gathercalling this kernel. It went from about 5.3 ms to about 1.5 ms, level with the hand-written branch-free loop.mask_gather_u32_underimplements. Mask-risc does not call it yet; that is a follow-up.This depends on the call site. In a standalone ndarray example, with the old loop inlined, the compiler had already removed the branch. There the old and new forms time the same (ratio 1.00–1.04), and a predictable all-set source did not speed up the old form either. The branch-free form removes the dependence on that compiler choice; it does not speed up every caller. The doc comment says so.
Tests
gather & gate. Covered: lengths 0/1/63/64/65/67/130, dirty output buffers, and gates (empty, full, random) with their tail bits pastnset.src_rows == 0with an emptysrcgives all-false and no panic, for both functions. A short gate buffer panics.simd-masking-parity: new check0xD01. All 14 groups are bit-identical on this AVX-512 host.src_rows == 0return, an emptysrcpanics.cargo clippy -p ndarray --all-targets -- -D warningsandcargo fmtare clean. The 165simd_masking_opslib tests and the 2 doctests pass.Not in this PR
MaskOp::Gather { under }, which would route the semijoin through the gated kernel. That change belongs in lance-graph.🤖 Generated with Claude Code
https://claude.ai/code/session_01HdJxpkHATkdNL2veorKao2
Generated by Claude Code
Summary by CodeRabbit