@@ -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