Skip to content

Commit d57ede4

Browse files
committed
simd: branch-free mask_gather_u32 + gated mask_gather_u32_under
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
1 parent 11089f6 commit d57ede4

3 files changed

Lines changed: 208 additions & 5 deletions

File tree

‎crates/simd-masking-parity/src/lib.rs‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ use ndarray::simd::{
3737
eq_u32_via_to_mask, eq_u64_to_mask, eq_u8_to_mask, ge_i32_to_mask, ge_i32_to_mask_under, ge_u64_to_mask,
3838
ge_u8_to_mask, gt_i32_to_mask, gt_i32_to_mask_under, gt_u64_to_mask, gt_u8_to_mask, le_i32_to_mask,
3939
le_i32_to_mask_under, le_u64_to_mask, le_u8_to_mask, lt_i32_to_mask, lt_i32_to_mask_under, lt_u64_to_mask,
40-
lt_u8_to_mask, mask_all, mask_and, mask_and_assign, mask_andnot, mask_andnot_assign, mask_any, mask_gather_u32,
40+
lt_u8_to_mask, mask_all, mask_and, mask_and_assign, mask_andnot, mask_andnot_assign, mask_any, mask_gather_u32, mask_gather_u32_under,
4141
mask_not, mask_not_assign, mask_or, mask_or_assign, mask_scatter_or_u32, mask_set_range, mask_shift_morton,
4242
mask_ternlog, mask_ternlog_any, mask_ternlog_assign, mask_ternlog_popcount, mask_xor, mask_xor_assign,
4343
masked_group_sum_i32, masked_group_sum_i32_via, masked_key_run_count_u32, masked_max_i32, masked_min_i32,
@@ -1269,6 +1269,20 @@ fn check_gather_scatter_group() -> Result<(), u32> {
12691269
return Err(0xD00);
12701270
}
12711271

1272+
// ── mask_gather_u32_under ────────────────────────────────────────
1273+
// Gate tail bits past `n` are left dirty on purpose: a kernel that
1274+
// used them would address `index` out of bounds.
1275+
let gate: Vec<u64> = (0..out_len).map(|_| rng.next()).collect();
1276+
let mut out = vec![u64::MAX; out_len];
1277+
mask_gather_u32_under(&src, src_rows, &index, &gate, &mut out);
1278+
let want = reference_mask(n, out_len, |i| {
1279+
let idx = index[i] as usize;
1280+
(gate[i / 64] >> (i % 64)) & 1 == 1 && idx < src_rows && (src[idx / 64] >> (idx % 64)) & 1 == 1
1281+
});
1282+
if out != want {
1283+
return Err(0xD01);
1284+
}
1285+
12721286
// ── mask_scatter_or_u32 ──────────────────────────────────────────
12731287
// `out2` starts explicitly zeroed (the caller's job now, not the
12741288
// primitive's): the accumulate contract means a dirty prefill would

‎src/simd.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -818,6 +818,7 @@ pub use crate::simd_masking_ops::{
818818
mask_andnot_assign,
819819
mask_any,
820820
mask_gather_u32,
821+
mask_gather_u32_under,
821822
mask_not,
822823
mask_not_assign,
823824
mask_or,

‎src/simd_masking_ops.rs‎

Lines changed: 192 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1076,6 +1076,20 @@ pub fn masked_strided_group_sum(
10761076
/// no input mask to skip zero words against — every entry of `index` must
10771077
/// be consulted regardless of what `src` contains.
10781078
///
1079+
/// **Branch-free per row.** Each row's bit is shifted into the output word
1080+
/// unconditionally; the out-of-range test becomes a select on the address
1081+
/// plus an AND on the bit (see `gather_bit`). The earlier form branched on
1082+
/// `idx < src_rows && bit == 1`. Whether that compiles to a branch depends
1083+
/// on the call site. Measured on a 1M-row lane with an unpredictable source:
1084+
/// called through lance-graph's mask-risc `Gather`, it ran 5.3 ms and the
1085+
/// branch-free form runs 1.5 ms, level with a hand-written branch-free loop
1086+
/// (`D-GATED-GATHER-0`). Inlined into a plain loop, the compiler had already
1087+
/// removed the branch, and the two forms time the same. The branch-free form
1088+
/// removes the dependence on that choice.
1089+
///
1090+
/// When the caller already holds the rows it wants — a conjunct's survivors —
1091+
/// use [`mask_gather_u32_under`], which reads `index` only at those rows.
1092+
///
10791093
/// # Panics
10801094
///
10811095
/// Panics if `out_words.len() < index.len().div_ceil(64)`, or if
@@ -1113,20 +1127,116 @@ pub fn mask_gather_u32(src: &[u64], src_rows: usize, index: &[u32], out_words: &
11131127
for w in out_words.iter_mut() {
11141128
*w = 0;
11151129
}
1130+
// With no source rows every index is out of range, and `src` may be
1131+
// empty, so `gather_bit`'s fallback read of `src[0]` would not exist.
1132+
if src_rows == 0 {
1133+
return;
1134+
}
11161135
for (w, out_word) in out_words.iter_mut().enumerate().take(words) {
11171136
let base = w * 64;
11181137
let live = (n - base).min(64);
11191138
let mut acc = 0u64;
11201139
for lane in 0..live {
1121-
let idx = index[base + lane] as usize;
1122-
if idx < src_rows && (src[idx / 64] >> (idx % 64)) & 1 == 1 {
1123-
acc |= 1u64 << lane;
1124-
}
1140+
acc |= gather_bit(src, src_rows, index[base + lane] as usize) << lane;
11251141
}
11261142
*out_word = acc;
11271143
}
11281144
}
11291145

1146+
/// Bit `idx` of `src`, or 0 when `idx >= src_rows` — without a branch on the
1147+
/// bit's value. An out-of-range address is redirected to row 0 by a select
1148+
/// and its result is cleared by `& ok`, so the read is always in bounds.
1149+
///
1150+
/// Requires `src_rows >= 1` (so `src[0]` exists); both gathers return early
1151+
/// on `src_rows == 0` before calling this.
1152+
#[inline(always)]
1153+
fn gather_bit(src: &[u64], src_rows: usize, idx: usize) -> u64 {
1154+
let ok = (idx < src_rows) as u64;
1155+
let j = if ok == 1 { idx } else { 0 };
1156+
((src[j / 64] >> (j % 64)) & 1) & ok
1157+
}
1158+
1159+
/// Gated gather: bit `i` of the result is `under[i] && src[index[i]]`, with
1160+
/// `index[i]` read **only** for rows whose `under` bit is set.
1161+
///
1162+
/// Same addressing and the same out-of-range-is-false contract as
1163+
/// [`mask_gather_u32`]; the result equals `mask_gather_u32(..) & under`. The
1164+
/// difference is the schedule: a zero `under` word costs one test, and inside
1165+
/// a live word only the set bits are visited (`trailing_zeros`), so the work
1166+
/// is proportional to `popcount(under)`, not to `index.len()`. This is the
1167+
/// semijoin a conjunction needs — `gate ∧ gather(fk, foreign)` — without
1168+
/// gathering the rows the gate has already rejected.
1169+
///
1170+
/// Measured (lance-graph `D-GATED-GATHER-0`, 1M rows): fastest or within
1171+
/// noise of fastest at every gate density from 0.01 % to 100 %, and at 100 %
1172+
/// no slower than the full branch-free gather. A word-granular gate (gather
1173+
/// all 64 lanes of every live word) is not a substitute: on scattered 1 %
1174+
/// survivors it issues ~47× the loads.
1175+
///
1176+
/// `under` bits at or past `index.len()` are ignored, never used to address
1177+
/// `index`. `out_words` is **fully overwritten**: every bit past
1178+
/// `index.len()` and every surplus word is written `0`.
1179+
///
1180+
/// # Panics
1181+
///
1182+
/// Panics if `out_words.len()` or `under.len()` is below
1183+
/// `index.len().div_ceil(64)`, or if `src.len() < src_rows.div_ceil(64)`.
1184+
///
1185+
/// # Examples
1186+
///
1187+
/// ```
1188+
/// use ndarray::simd::{mask_gather_u32, mask_gather_u32_under};
1189+
///
1190+
/// let src = [0b101u64]; // rows 0 and 2 of a 3-row source are set
1191+
/// let index = [2u32, 1, 0, 5];
1192+
/// let mut full = [0u64; 1];
1193+
/// mask_gather_u32(&src, 3, &index, &mut full);
1194+
/// assert_eq!(full[0], 0b0101);
1195+
/// // keep only rows 2 and 3 of `index`: row 2 reads src row 0 (set),
1196+
/// // row 3 is out of range (false)
1197+
/// let mut gated = [u64::MAX; 1];
1198+
/// mask_gather_u32_under(&src, 3, &index, &[0b1100], &mut gated);
1199+
/// assert_eq!(gated[0], 0b0100);
1200+
/// ```
1201+
#[inline]
1202+
pub fn mask_gather_u32_under(src: &[u64], src_rows: usize, index: &[u32], under: &[u64], out_words: &mut [u64]) {
1203+
let n = index.len();
1204+
let words = mask_words_for(n);
1205+
assert!(
1206+
out_words.len() >= words,
1207+
"mask_gather_u32_under: out_words.len()={} < required {}",
1208+
out_words.len(),
1209+
words
1210+
);
1211+
assert!(under.len() >= words, "mask_gather_u32_under: under.len()={} < required {}", under.len(), words);
1212+
let src_words = mask_words_for(src_rows);
1213+
assert!(src.len() >= src_words, "mask_gather_u32_under: src.len()={} < required {}", src.len(), src_words);
1214+
1215+
for w in out_words.iter_mut() {
1216+
*w = 0;
1217+
}
1218+
if src_rows == 0 {
1219+
return;
1220+
}
1221+
for w in 0..words {
1222+
let base = w * 64;
1223+
let mut bits = under[w];
1224+
// A gate bit past `index.len()` names no row: drop it before it can
1225+
// address `index` out of bounds.
1226+
let valid = n - base;
1227+
if valid < 64 {
1228+
bits &= (1u64 << valid) - 1;
1229+
}
1230+
let mut acc = 0u64;
1231+
while bits != 0 {
1232+
let lane = bits.trailing_zeros() as usize;
1233+
bits &= bits - 1;
1234+
acc |= gather_bit(src, src_rows, index[base + lane] as usize) << lane;
1235+
}
1236+
out_words[w] = acc;
1237+
}
1238+
}
1239+
11301240
/// Scatter-OR by row index: for every element `i` selected by `src` (bit `i`
11311241
/// set, `i < index.len()`), sets bit `index[i]` of the result — provided
11321242
/// `index[i] < out_rows`.
@@ -6246,6 +6356,84 @@ mod tests {
62466356
mask_gather_u32(&src, 65, &index, &mut out); // claims 65 rows
62476357
}
62486358

6359+
#[test]
6360+
fn mask_gather_u32_zero_src_rows_is_all_false_even_with_empty_src() {
6361+
// `src_rows == 0` needs no source words, so `src` may be empty: the
6362+
// branch-free fallback read of `src[0]` must not happen.
6363+
let index = [0u32, 1, 7];
6364+
let mut out = [u64::MAX; 1];
6365+
mask_gather_u32(&[], 0, &index, &mut out);
6366+
assert_eq!(out[0], 0);
6367+
let mut gated = [u64::MAX; 1];
6368+
mask_gather_u32_under(&[], 0, &index, &[u64::MAX], &mut gated);
6369+
assert_eq!(gated[0], 0);
6370+
}
6371+
6372+
#[test]
6373+
fn mask_gather_u32_under_matches_gather_and_gate_across_the_tail() {
6374+
for &n in &[0usize, 1, 63, 64, 65, 67, 130] {
6375+
let mut seed = 0x2468_ACE0_1357_9BDFu64;
6376+
let src_rows = 40usize;
6377+
let src: Vec<u64> = (0..src_rows.div_ceil(64))
6378+
.map(|_| splitmix(&mut seed))
6379+
.collect();
6380+
let index: Vec<u32> = (0..n)
6381+
.map(|i| {
6382+
if i % 3 == 0 {
6383+
(src_rows as u64 + 5 + i as u64) as u32
6384+
} else {
6385+
(splitmix(&mut seed) % src_rows as u64) as u32
6386+
}
6387+
})
6388+
.collect();
6389+
let out_words = n.div_ceil(64).max(1);
6390+
// Gates: empty, full, random — each with its tail bits past `n`
6391+
// deliberately set, so a missing clamp reads `index` out of range.
6392+
let gates: [Vec<u64>; 3] = [
6393+
vec![0u64; out_words],
6394+
vec![u64::MAX; out_words],
6395+
(0..out_words).map(|_| splitmix(&mut seed)).collect(),
6396+
];
6397+
for gate in &gates {
6398+
let want: Vec<bool> = naive_gather(&src, src_rows, &index)
6399+
.iter()
6400+
.enumerate()
6401+
.map(|(i, &b)| b && (gate[i / 64] >> (i % 64)) & 1 == 1)
6402+
.collect();
6403+
let want = bits_to_words(&want);
6404+
let mut out = vec![u64::MAX; out_words + 1];
6405+
mask_gather_u32_under(&src, src_rows, &index, gate, &mut out);
6406+
assert_eq!(&out[..want.len()], &want[..], "gated gather mismatch at n={n}");
6407+
assert_eq!(out[out_words], 0, "surplus word must be cleared at n={n}");
6408+
}
6409+
}
6410+
}
6411+
6412+
#[test]
6413+
fn mask_gather_u32_under_keeps_only_gated_rows() {
6414+
// Every row's address selects a set source bit, so the gate alone
6415+
// decides the result. (Whether ungated rows are READ is not
6416+
// observable through this API; the dirty-tail gates in the parity
6417+
// test above are what prove a gate bit past `index.len()` never
6418+
// addresses `index`.)
6419+
let src = [u64::MAX]; // 64 rows, all set
6420+
let index = [0u32, 1, 2, 3];
6421+
let mut out = [0u64; 1];
6422+
mask_gather_u32_under(&src, 64, &index, &[0b0100], &mut out);
6423+
assert_eq!(out[0], 0b0100, "only the gated row survives");
6424+
mask_gather_u32_under(&src, 64, &index, &[0], &mut out);
6425+
assert_eq!(out[0], 0, "an empty gate gathers nothing");
6426+
}
6427+
6428+
#[test]
6429+
#[should_panic(expected = "under.len()")]
6430+
fn mask_gather_u32_under_rejects_short_gate() {
6431+
let src = [0u64; 1];
6432+
let index = vec![0u32; 65];
6433+
let mut out = [0u64; 2];
6434+
mask_gather_u32_under(&src, 1, &index, &[u64::MAX], &mut out);
6435+
}
6436+
62496437
fn naive_scatter(src_bits: &[bool], index: &[u32], out_rows: usize) -> Vec<bool> {
62506438
let mut out = vec![false; out_rows];
62516439
for (i, &selected) in src_bits.iter().enumerate() {

0 commit comments

Comments
 (0)