Skip to main content

vibesql_executor/select/columnar/simd_ops/
mod.rs

1//! Auto-vectorized SIMD operations for columnar data processing.
2//!
3//! # Performance Architecture
4//!
5//! These functions are structured to enable LLVM auto-vectorization. They achieve
6//! equivalent performance to explicit SIMD (e.g., the `wide` crate) without the
7//! complexity of platform-specific code.
8//!
9//! ## Why This Pattern Works
10//!
11//! LLVM can auto-vectorize loops when:
12//! 1. Loop bounds are known or predictable
13//! 2. Memory access is sequential
14//! 3. Operations are independent across lanes
15//!
16//! The 4-accumulator pattern breaks loop-carried dependencies, allowing LLVM to
17//! use SIMD registers effectively:
18//!
19//! ```text
20//! // BAD: Single accumulator creates dependency chain
21//! for x in data { sum += x; }  // Each add waits for previous
22//!
23//! // GOOD: Four accumulators enable parallel execution
24//! for chunk in data.chunks(4) {
25//!     s0 += chunk[0];  // These four adds can execute
26//!     s1 += chunk[1];  // simultaneously in SIMD lanes
27//!     s2 += chunk[2];
28//!     s3 += chunk[3];
29//! }
30//! ```
31//!
32//! ## Benchmark Results (10M elements, Apple Silicon)
33//!
34//! | Operation | wide crate | auto-vectorized | naive iter |
35//! |-----------|------------|-----------------|------------|
36//! | sum_f64   | 2.0 ms     | 2.0 ms (1.0x)   | 7.8 ms     |
37//! | min_f64   | 1.5 ms     | 1.5 ms (1.0x)   | 1.4 ms     |
38//!
39//! ## WARNING
40//!
41//! DO NOT "simplify" these functions to use `.iter().sum()` or similar patterns.
42//! While cleaner-looking, they can be 3-4x slower due to floating-point
43//! associativity constraints preventing vectorization.
44//!
45//! If you need to modify these functions, run the SIMD benchmark first:
46//! ```bash
47//! cargo bench --bench tpch -- Q6
48//! ```
49
50mod aggregate;
51mod comparison;
52pub mod intrinsics;
53mod logical;
54
55pub use aggregate::*;
56pub use comparison::*;
57pub use logical::*;
58
59// ============================================================================
60// PACKED BITMASK TYPE
61// ============================================================================
62// Using packed bitmasks (1 bit per row) instead of Vec<bool> (1 byte per row)
63// provides 8x memory reduction and enables native SIMD bitwise operations.
64
65/// A packed bitmask where each bit represents whether a row passes a filter.
66///
67/// # Memory Efficiency
68///
69/// For 6M rows (lineitem at TPC-H SF 1):
70/// - `Vec<bool>`: 6 MB per predicate
71/// - `PackedMask`: 750 KB per predicate (8x reduction)
72///
73/// # Performance Benefits
74///
75/// - Bitwise AND is a native SIMD instruction (`vpand` on x86, `and` on ARM)
76/// - Better cache utilization due to reduced memory footprint
77/// - Faster popcount using `count_ones()` intrinsic
78///
79/// # Implementation
80///
81/// Bits are stored in little-endian order within each u64 word:
82/// - Bit 0 of word 0 = row 0
83/// - Bit 63 of word 0 = row 63
84/// - Bit 0 of word 1 = row 64
85/// - etc.
86#[derive(Debug, Clone, PartialEq, Eq)]
87pub struct PackedMask {
88    /// The packed bits. Each u64 holds 64 row flags.
89    words: Vec<u64>,
90    /// The actual number of rows (since the last word may be partial).
91    len: usize,
92}
93
94impl PackedMask {
95    /// Creates a new mask with all bits set (all rows pass).
96    #[inline]
97    pub fn new_all_set(len: usize) -> Self {
98        if len == 0 {
99            return Self { words: vec![], len: 0 };
100        }
101        let num_words = len.div_ceil(64);
102        let mut words = vec![u64::MAX; num_words];
103
104        // Clear unused bits in the last word
105        let remainder = len % 64;
106        if remainder != 0 {
107            words[num_words - 1] = (1u64 << remainder) - 1;
108        }
109
110        Self { words, len }
111    }
112
113    /// Creates a new mask with all bits clear (no rows pass).
114    #[inline]
115    pub fn new_all_clear(len: usize) -> Self {
116        if len == 0 {
117            return Self { words: vec![], len: 0 };
118        }
119        let num_words = len.div_ceil(64);
120        Self { words: vec![0; num_words], len }
121    }
122
123    /// Returns the number of rows this mask represents.
124    #[inline]
125    pub fn len(&self) -> usize {
126        self.len
127    }
128
129    /// Returns true if this mask is empty.
130    #[inline]
131    pub fn is_empty(&self) -> bool {
132        self.len == 0
133    }
134
135    /// Gets the value of the bit at the given index.
136    #[inline]
137    pub fn get(&self, index: usize) -> bool {
138        debug_assert!(index < self.len, "index out of bounds");
139        let word_idx = index / 64;
140        let bit_idx = index % 64;
141        (self.words[word_idx] >> bit_idx) & 1 == 1
142    }
143
144    /// Sets the bit at the given index.
145    #[inline]
146    pub fn set(&mut self, index: usize, value: bool) {
147        debug_assert!(index < self.len, "index out of bounds");
148        let word_idx = index / 64;
149        let bit_idx = index % 64;
150        if value {
151            self.words[word_idx] |= 1u64 << bit_idx;
152        } else {
153            self.words[word_idx] &= !(1u64 << bit_idx);
154        }
155    }
156
157    /// Performs in-place bitwise AND with another mask.
158    ///
159    /// This is the key operation for combining multiple predicates.
160    /// Uses native bitwise AND which maps directly to SIMD instructions.
161    #[inline]
162    pub fn and_inplace(&mut self, other: &PackedMask) {
163        debug_assert_eq!(self.len, other.len, "mask lengths must match");
164
165        // Process in chunks of 4 for potential auto-vectorization
166        let chunks = self.words.len() / 4;
167        for i in 0..chunks {
168            let base = i * 4;
169            self.words[base] &= other.words[base];
170            self.words[base + 1] &= other.words[base + 1];
171            self.words[base + 2] &= other.words[base + 2];
172            self.words[base + 3] &= other.words[base + 3];
173        }
174
175        // Handle remainder
176        for i in (chunks * 4)..self.words.len() {
177            self.words[i] &= other.words[i];
178        }
179    }
180
181    /// Performs in-place bitwise OR with another mask.
182    ///
183    /// This is used for combining multiple predicates in IN lists.
184    /// Uses native bitwise OR which maps directly to SIMD instructions.
185    #[inline]
186    pub fn or_inplace(&mut self, other: &PackedMask) {
187        debug_assert_eq!(self.len, other.len, "mask lengths must match");
188
189        // Process in chunks of 4 for potential auto-vectorization
190        let chunks = self.words.len() / 4;
191        for i in 0..chunks {
192            let base = i * 4;
193            self.words[base] |= other.words[base];
194            self.words[base + 1] |= other.words[base + 1];
195            self.words[base + 2] |= other.words[base + 2];
196            self.words[base + 3] |= other.words[base + 3];
197        }
198
199        // Handle remainder
200        for i in (chunks * 4)..self.words.len() {
201            self.words[i] |= other.words[i];
202        }
203    }
204
205    /// Returns a new mask with all bits inverted.
206    ///
207    /// This is used for NOT IN operations.
208    #[inline]
209    pub fn not(&self) -> PackedMask {
210        let mut result = PackedMask { words: vec![0; self.words.len()], len: self.len };
211
212        // Process in chunks of 4 for potential auto-vectorization
213        let chunks = self.words.len() / 4;
214        for i in 0..chunks {
215            let base = i * 4;
216            result.words[base] = !self.words[base];
217            result.words[base + 1] = !self.words[base + 1];
218            result.words[base + 2] = !self.words[base + 2];
219            result.words[base + 3] = !self.words[base + 3];
220        }
221
222        // Handle remainder
223        for i in (chunks * 4)..self.words.len() {
224            result.words[i] = !self.words[i];
225        }
226
227        // Clear unused bits in the last word
228        let remainder = self.len % 64;
229        let num_words = result.words.len();
230        if remainder != 0 && num_words > 0 {
231            result.words[num_words - 1] &= (1u64 << remainder) - 1;
232        }
233
234        result
235    }
236
237    /// Counts the number of set bits (rows that pass the filter).
238    ///
239    /// Uses the `count_ones()` intrinsic which maps to hardware popcount.
240    #[inline]
241    pub fn count_ones(&self) -> usize {
242        // Use 4-accumulator pattern for better throughput
243        let chunks = self.words.len() / 4;
244        let (mut c0, mut c1, mut c2, mut c3) = (0u32, 0u32, 0u32, 0u32);
245
246        for i in 0..chunks {
247            let base = i * 4;
248            c0 += self.words[base].count_ones();
249            c1 += self.words[base + 1].count_ones();
250            c2 += self.words[base + 2].count_ones();
251            c3 += self.words[base + 3].count_ones();
252        }
253
254        let mut total = (c0 + c1 + c2 + c3) as usize;
255        for i in (chunks * 4)..self.words.len() {
256            total += self.words[i].count_ones() as usize;
257        }
258        total
259    }
260
261    /// Creates a packed mask from a boolean slice.
262    ///
263    /// Used for compatibility with existing code during migration.
264    #[inline]
265    pub fn from_bool_slice(slice: &[bool]) -> Self {
266        let len = slice.len();
267        if len == 0 {
268            return Self { words: vec![], len: 0 };
269        }
270
271        let num_words = len.div_ceil(64);
272        let mut words = vec![0u64; num_words];
273
274        // Process 64 bits at a time
275        for (word_idx, word) in words.iter_mut().enumerate() {
276            let start = word_idx * 64;
277            let end = (start + 64).min(len);
278            let mut bits = 0u64;
279            for (bit_idx, &value) in slice[start..end].iter().enumerate() {
280                if value {
281                    bits |= 1u64 << bit_idx;
282                }
283            }
284            *word = bits;
285        }
286
287        Self { words, len }
288    }
289
290    /// Converts this packed mask back to a boolean vector.
291    ///
292    /// Used for compatibility with existing code during migration.
293    #[inline]
294    pub fn to_bool_vec(&self) -> Vec<bool> {
295        let mut result = Vec::with_capacity(self.len);
296        for i in 0..self.len {
297            result.push(self.get(i));
298        }
299        result
300    }
301
302    /// Returns a reference to the underlying word slice.
303    #[inline]
304    pub fn words(&self) -> &[u64] {
305        &self.words
306    }
307
308    /// Creates a PackedMask from raw words and length.
309    ///
310    /// This is used internally by comparison functions that build masks.
311    /// The caller must ensure the words are valid for the given length.
312    #[inline]
313    pub(crate) fn from_words(words: Vec<u64>, len: usize) -> Self {
314        Self { words, len }
315    }
316}
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321
322    #[test]
323    fn test_packed_mask_new_all_set() {
324        let mask = PackedMask::new_all_set(100);
325        assert_eq!(mask.len(), 100);
326        for i in 0..100 {
327            assert!(mask.get(i), "bit {} should be set", i);
328        }
329
330        // Empty mask
331        let empty = PackedMask::new_all_set(0);
332        assert_eq!(empty.len(), 0);
333        assert!(empty.is_empty());
334    }
335
336    #[test]
337    fn test_packed_mask_new_all_clear() {
338        let mask = PackedMask::new_all_clear(100);
339        assert_eq!(mask.len(), 100);
340        for i in 0..100 {
341            assert!(!mask.get(i), "bit {} should be clear", i);
342        }
343    }
344
345    #[test]
346    fn test_packed_mask_set_get() {
347        let mut mask = PackedMask::new_all_clear(128);
348
349        // Set some bits
350        mask.set(0, true);
351        mask.set(63, true);
352        mask.set(64, true);
353        mask.set(127, true);
354
355        assert!(mask.get(0));
356        assert!(mask.get(63));
357        assert!(mask.get(64));
358        assert!(mask.get(127));
359        assert!(!mask.get(1));
360        assert!(!mask.get(62));
361        assert!(!mask.get(65));
362    }
363
364    #[test]
365    fn test_packed_mask_and_inplace() {
366        let mut mask_a = PackedMask::new_all_set(128);
367        let mut mask_b = PackedMask::new_all_clear(128);
368
369        // Set every other bit in mask_b
370        for i in (0..128).step_by(2) {
371            mask_b.set(i, true);
372        }
373
374        mask_a.and_inplace(&mask_b);
375
376        for i in 0..128 {
377            if i % 2 == 0 {
378                assert!(mask_a.get(i), "bit {} should be set (even)", i);
379            } else {
380                assert!(!mask_a.get(i), "bit {} should be clear (odd)", i);
381            }
382        }
383    }
384
385    #[test]
386    fn test_packed_mask_count_ones() {
387        let mask = PackedMask::new_all_set(100);
388        assert_eq!(mask.count_ones(), 100);
389
390        let mask = PackedMask::new_all_clear(100);
391        assert_eq!(mask.count_ones(), 0);
392
393        let mut mask = PackedMask::new_all_clear(128);
394        for i in (0..128).step_by(2) {
395            mask.set(i, true);
396        }
397        assert_eq!(mask.count_ones(), 64);
398    }
399
400    #[test]
401    fn test_packed_mask_from_bool_slice() {
402        let bools = vec![true, false, true, true, false, true];
403        let mask = PackedMask::from_bool_slice(&bools);
404
405        assert_eq!(mask.len(), 6);
406        assert!(mask.get(0));
407        assert!(!mask.get(1));
408        assert!(mask.get(2));
409        assert!(mask.get(3));
410        assert!(!mask.get(4));
411        assert!(mask.get(5));
412        assert_eq!(mask.count_ones(), 4);
413    }
414
415    #[test]
416    fn test_packed_mask_to_bool_vec() {
417        let mut mask = PackedMask::new_all_clear(6);
418        mask.set(0, true);
419        mask.set(2, true);
420        mask.set(3, true);
421        mask.set(5, true);
422
423        let bools = mask.to_bool_vec();
424        assert_eq!(bools, vec![true, false, true, true, false, true]);
425    }
426
427    #[test]
428    fn test_packed_mask_round_trip() {
429        let original = vec![true, false, true, true, false, true, false, false, true];
430        let mask = PackedMask::from_bool_slice(&original);
431        let result = mask.to_bool_vec();
432        assert_eq!(original, result);
433    }
434
435    #[test]
436    fn test_packed_mask_partial_word() {
437        // Test with sizes that don't fill a complete u64 word
438        for size in [1, 7, 15, 63, 65, 100, 127, 129] {
439            let mask = PackedMask::new_all_set(size);
440            assert_eq!(mask.count_ones(), size, "failed for size {}", size);
441
442            let mask = PackedMask::new_all_clear(size);
443            assert_eq!(mask.count_ones(), 0, "failed for size {}", size);
444
445            // Test conversion round-trip
446            let bools = vec![true; size];
447            let mask = PackedMask::from_bool_slice(&bools);
448            assert_eq!(mask.count_ones(), size, "round-trip failed for size {}", size);
449            assert_eq!(mask.to_bool_vec(), bools, "round-trip failed for size {}", size);
450        }
451    }
452}