Skip to main content

polars_utils/
clmul.rs

1#[cfg(all(target_arch = "x86_64", target_feature = "pclmulqdq"))]
2#[inline]
3fn intel_clmul64(x: u64, y: u64) -> u64 {
4    use core::arch::x86_64::*;
5    unsafe {
6        // SAFETY: we have the target feature.
7        _mm_cvtsi128_si64(_mm_clmulepi64_si128(
8            _mm_cvtsi64_si128(x as i64),
9            _mm_cvtsi64_si128(y as i64),
10            0,
11        )) as u64
12    }
13}
14
15#[cfg(all(
16    target_arch = "aarch64",
17    target_feature = "neon",
18    target_feature = "aes"
19))]
20#[inline]
21fn arm_clmul64(x: u64, y: u64) -> u64 {
22    unsafe {
23        // SAFETY: we have the target feature.
24        use core::arch::aarch64::*;
25        vmull_p64(x, y) as u64
26    }
27}
28
29#[inline]
30pub fn portable_clmul64(x: u64, mut y: u64) -> u64 {
31    let mut out = 0;
32    while y > 0 {
33        let lsb = y.isolate_lowest_one();
34        out ^= x.wrapping_mul(lsb);
35        y ^= lsb;
36    }
37    out
38}
39
40// Computes the carryless multiplication of x and y.
41#[inline]
42pub fn clmul64(x: u64, y: u64) -> u64 {
43    #[cfg(all(target_arch = "x86_64", target_feature = "pclmulqdq"))]
44    return intel_clmul64(x, y);
45
46    #[cfg(all(
47        target_arch = "aarch64",
48        target_feature = "neon",
49        target_feature = "aes"
50    ))]
51    return arm_clmul64(x, y);
52
53    #[allow(unreachable_code)]
54    portable_clmul64(x, y)
55}
56
57#[inline]
58pub fn portable_prefix_xorsum(x: u64) -> u64 {
59    portable_prefix_xorsum_inclusive(x << 1)
60}
61
62// Computes for each bit i the XOR of bits[0..i].
63#[inline]
64pub fn prefix_xorsum(x: u64) -> u64 {
65    #[cfg(all(target_arch = "x86_64", target_feature = "pclmulqdq"))]
66    return intel_clmul64(x, u64::MAX ^ 1);
67
68    #[cfg(all(
69        target_arch = "aarch64",
70        target_feature = "neon",
71        target_feature = "aes"
72    ))]
73    return arm_clmul64(x, u64::MAX ^ 1);
74
75    #[allow(unreachable_code)]
76    portable_prefix_xorsum(x)
77}
78
79#[inline]
80pub fn portable_prefix_xorsum_inclusive(mut x: u64) -> u64 {
81    for i in 0..6 {
82        x ^= x << (1 << i);
83    }
84    x
85}
86
87// Computes for each bit i the XOR of bits[0..=i].
88#[inline]
89pub fn prefix_xorsum_inclusive(x: u64) -> u64 {
90    #[cfg(all(target_arch = "x86_64", target_feature = "pclmulqdq"))]
91    return intel_clmul64(x, u64::MAX);
92
93    #[cfg(all(
94        target_arch = "aarch64",
95        target_feature = "neon",
96        target_feature = "aes"
97    ))]
98    return arm_clmul64(x, u64::MAX);
99
100    #[allow(unreachable_code)]
101    portable_prefix_xorsum_inclusive(x)
102}
103
104#[cfg(test)]
105mod test {
106    use rand::prelude::*;
107
108    use super::*;
109
110    #[test]
111    fn test_clmul() {
112        // Verify platform-specific clmul to portable.
113        let mut rng = StdRng::seed_from_u64(0xdeadbeef);
114        for _ in 0..100 {
115            let x = rng.random();
116            let y = rng.random();
117            assert_eq!(portable_clmul64(x, y), clmul64(x, y));
118        }
119
120        // Verify portable clmul for known test vectors.
121        assert_eq!(
122            portable_clmul64(0x8b44729195dde0ef, 0xb976c5ae2726fab0),
123            0x4ae14eae84899290
124        );
125        assert_eq!(
126            portable_clmul64(0x399b6ed00c44b301, 0x693341db5acb2ff0),
127            0x48dfa88344823ff0
128        );
129        assert_eq!(
130            portable_clmul64(0xdf4c9f6e60deb640, 0x6d4bcdb217ac4880),
131            0x7300ffe474792000
132        );
133        assert_eq!(
134            portable_clmul64(0xa7adf3c53a200a51, 0x818cb40fe11b431e),
135            0x6a280181d521797e
136        );
137        assert_eq!(
138            portable_clmul64(0x5e78e12b744f228c, 0x4225ff19e9273266),
139            0xa48b73cafb9665a8
140        );
141    }
142
143    #[test]
144    fn test_prefix_xorsum() {
145        // Verify platform-specific prefix_xorsum to portable.
146        let mut rng = StdRng::seed_from_u64(0xdeadbeef);
147        for _ in 0..100 {
148            let x = rng.random();
149            assert_eq!(portable_prefix_xorsum(x), prefix_xorsum(x));
150        }
151
152        // Verify portable prefix_xorsum for known test vectors.
153        assert_eq!(
154            portable_prefix_xorsum(0x8b44729195dde0ef),
155            0x0d87a31ee696bf4a
156        );
157        assert_eq!(
158            portable_prefix_xorsum(0xb976c5ae2726fab0),
159            0x2e5b79343a3b5320
160        );
161        assert_eq!(
162            portable_prefix_xorsum(0x399b6ed00c44b301),
163            0xd1124b600878ddfe
164        );
165        assert_eq!(
166            portable_prefix_xorsum(0x693341db5acb2ff0),
167            0x4e227e926c8dcaa0
168        );
169        assert_eq!(
170            portable_prefix_xorsum(0xdf4c9f6e60deb640),
171            0x6a7715b44094db80
172        );
173    }
174}