Skip to main content

fxcp_core/operations/
simd.rs

1// SPDX-License-Identifier: GPL-2.0-or-later
2// Copyright (C) 2025 Joel Wirāmu Pauling <aenertia@aenertia.net>
3//
4//! SIMD zero-block detection (AVX-512, AVX2, NEON, fallback).
5
6#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
7use std::arch::x86_64::*;
8
9// ---------------------------------------------------------------------------
10// Public SIMD-accelerated zero-block detection
11// ---------------------------------------------------------------------------
12
13/// Check if a buffer is entirely zero. Uses architecture-specific SIMD:
14/// - x86_64: AVX-512 (256B/iter), AVX2 (32B/iter), runtime-detected
15/// - AArch64: NEON vmaxvq_u8 (64B/iter)
16/// - Generic: u128-aligned comparison fallback
17#[inline(always)]
18pub(crate) fn is_zero_block(buf: &[u8]) -> bool {
19    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
20    {
21        #[cfg(feature = "nightly")]
22        {
23            if is_x86_feature_detected!("avx512f") {
24                // SAFETY: AVX-512 support verified by is_x86_feature_detected! above.
25                return unsafe { is_zero_avx512(buf) };
26            }
27        }
28        if is_x86_feature_detected!("avx2") {
29            // SAFETY: AVX2 support verified by is_x86_feature_detected! above.
30            return unsafe { is_zero_avx2(buf) };
31        }
32    }
33
34    #[cfg(target_arch = "aarch64")]
35    {
36        return is_zero_neon(buf);
37    }
38
39    #[allow(unreachable_code)]
40    {
41        // SAFETY: u128 has no invalid bit patterns, so align_to is safe for
42        // any byte slice. We check all three segments (prefix, chunks, suffix).
43        let (prefix, chunks, suffix) = unsafe { buf.align_to::<u128>() };
44        chunks.iter().all(|&x| x == 0) && prefix.iter().all(|&x| x == 0) && suffix.iter().all(|&x| x == 0)
45    }
46}
47
48#[cfg(target_arch = "aarch64")]
49fn is_zero_neon(buf: &[u8]) -> bool {
50    use std::arch::aarch64::*;
51    let len = buf.len();
52    let ptr = buf.as_ptr();
53    let mut i = 0;
54    // SAFETY: buf.as_ptr() is valid for buf.len() bytes. Loop bounds ensure
55    // we never read past the end. NEON intrinsics require valid aligned pointers,
56    // and vld1q_u8 handles unaligned loads.
57    unsafe {
58        while i + 64 <= len {
59            let a = vld1q_u8(ptr.add(i));
60            let b = vld1q_u8(ptr.add(i + 16));
61            let c = vld1q_u8(ptr.add(i + 32));
62            let d = vld1q_u8(ptr.add(i + 48));
63            let combined = vorrq_u8(vorrq_u8(a, b), vorrq_u8(c, d));
64            if vmaxvq_u8(combined) != 0 { return false; }
65            i += 64;
66        }
67    }
68    buf[i..].iter().all(|&b| b == 0)
69}
70
71#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), feature = "nightly"))]
72#[target_feature(enable = "avx512f")]
73unsafe fn is_zero_avx512(buf: &[u8]) -> bool {
74    let len = buf.len();
75    let ptr = buf.as_ptr();
76    let mut i = 0;
77    // SAFETY: caller guarantees AVX-512 is available (target_feature attribute).
78    // Loop bounds ensure all loads stay within buf. _mm512_loadu_si512 handles
79    // unaligned addresses.
80    unsafe {
81        while i + 256 <= len {
82            let a = _mm512_loadu_si512(ptr.add(i) as *const _);
83            let b = _mm512_loadu_si512(ptr.add(i + 64) as *const _);
84            let c = _mm512_loadu_si512(ptr.add(i + 128) as *const _);
85            let d = _mm512_loadu_si512(ptr.add(i + 192) as *const _);
86            let combined = _mm512_or_si512(_mm512_or_si512(a, b), _mm512_or_si512(c, d));
87            if _mm512_test_epi64_mask(combined, combined) != 0 { return false; }
88            i += 256;
89        }
90        while i + 64 <= len {
91            let a = _mm512_loadu_si512(ptr.add(i) as *const _);
92            if _mm512_test_epi64_mask(a, a) != 0 { return false; }
93            i += 64;
94        }
95        buf[i..].iter().all(|&b| b == 0)
96    }
97}
98
99#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
100#[target_feature(enable = "avx2")]
101unsafe fn is_zero_avx2(buf: &[u8]) -> bool {
102    let len = buf.len();
103    let ptr = buf.as_ptr();
104    let mut i = 0;
105    // SAFETY: caller guarantees AVX2 is available (target_feature attribute).
106    // Loop bounds ensure all loads stay within buf. _mm256_loadu_si256 handles
107    // unaligned addresses.
108    unsafe {
109        while i + 32 <= len {
110            let a = _mm256_loadu_si256(ptr.add(i) as *const _);
111            if _mm256_testz_si256(a, a) == 0 { return false; }
112            i += 32;
113        }
114        buf[i..].iter().all(|&b| b == 0)
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
121    use super::*;
122
123    #[test]
124    fn test_empty_buffer_is_zero() {
125        assert!(is_zero_block(&[]));
126    }
127
128    #[test]
129    fn test_single_zero_byte() {
130        assert!(is_zero_block(&[0]));
131    }
132
133    #[test]
134    fn test_single_nonzero_byte() {
135        assert!(!is_zero_block(&[1]));
136    }
137
138    #[test]
139    fn test_all_zeros_small() {
140        assert!(is_zero_block(&[0u8; 7]));
141    }
142
143    #[test]
144    fn test_all_zeros_16_bytes() {
145        assert!(is_zero_block(&[0u8; 16]));
146    }
147
148    #[test]
149    fn test_all_zeros_32_bytes() {
150        assert!(is_zero_block(&[0u8; 32]));
151    }
152
153    #[test]
154    fn test_all_zeros_64_bytes() {
155        assert!(is_zero_block(&[0u8; 64]));
156    }
157
158    #[test]
159    fn test_all_zeros_256_bytes() {
160        assert!(is_zero_block(&[0u8; 256]));
161    }
162
163    #[test]
164    fn test_all_zeros_4096_bytes() {
165        assert!(is_zero_block(&[0u8; 4096]));
166    }
167
168    #[test]
169    fn test_all_nonzero() {
170        assert!(!is_zero_block(&[0xFF; 4096]));
171    }
172
173    #[test]
174    fn test_nonzero_first_byte() {
175        let mut buf = vec![0u8; 4096];
176        buf[0] = 1;
177        assert!(!is_zero_block(&buf));
178    }
179
180    #[test]
181    fn test_nonzero_last_byte() {
182        let mut buf = vec![0u8; 4096];
183        buf[4095] = 1;
184        assert!(!is_zero_block(&buf));
185    }
186
187    #[test]
188    fn test_nonzero_middle_byte() {
189        let mut buf = vec![0u8; 4096];
190        buf[2048] = 0x42;
191        assert!(!is_zero_block(&buf));
192    }
193
194    #[test]
195    fn test_unaligned_sizes() {
196        for size in [1, 3, 15, 17, 31, 33, 63, 65, 127, 255, 257, 511, 1023] {
197            assert!(is_zero_block(&vec![0u8; size]), "failed for zero buf size={}", size);
198            let mut buf = vec![0u8; size];
199            buf[size - 1] = 1;
200            assert!(!is_zero_block(&buf), "failed for nonzero buf size={}", size);
201        }
202    }
203
204    #[test]
205    fn test_large_zero_block() {
206        assert!(is_zero_block(&vec![0u8; 1024 * 1024]));
207    }
208
209    #[test]
210    fn test_large_block_nonzero_at_end() {
211        let mut buf = vec![0u8; 1024 * 1024];
212        buf[1024 * 1024 - 1] = 1;
213        assert!(!is_zero_block(&buf));
214    }
215}