blake3/
ffi_avx512.rs

1use crate::{CVWords, IncrementCounter, BLOCK_LEN, OUT_LEN};
2
3// Unsafe because this may only be called on platforms supporting AVX-512.
4pub unsafe fn compress_in_place(
5    cv: &mut CVWords,
6    block: &[u8; BLOCK_LEN],
7    block_len: u8,
8    counter: u64,
9    flags: u8,
10) {
11    unsafe {
12        ffi::blake3_compress_in_place_avx512(
13            cv.as_mut_ptr(),
14            block.as_ptr(),
15            block_len,
16            counter,
17            flags,
18        )
19    }
20}
21
22// Unsafe because this may only be called on platforms supporting AVX-512.
23pub unsafe fn compress_xof(
24    cv: &CVWords,
25    block: &[u8; BLOCK_LEN],
26    block_len: u8,
27    counter: u64,
28    flags: u8,
29) -> [u8; 64] {
30    unsafe {
31        let mut out = [0u8; 64];
32        ffi::blake3_compress_xof_avx512(
33            cv.as_ptr(),
34            block.as_ptr(),
35            block_len,
36            counter,
37            flags,
38            out.as_mut_ptr(),
39        );
40        out
41    }
42}
43
44// Unsafe because this may only be called on platforms supporting AVX-512.
45pub unsafe fn hash_many<const N: usize>(
46    inputs: &[&[u8; N]],
47    key: &CVWords,
48    counter: u64,
49    increment_counter: IncrementCounter,
50    flags: u8,
51    flags_start: u8,
52    flags_end: u8,
53    out: &mut [u8],
54) {
55    unsafe {
56        // The Rust hash_many implementations do bounds checking on the `out`
57        // array, but the C implementations don't. Even though this is an unsafe
58        // function, assert the bounds here.
59        assert!(out.len() >= inputs.len() * OUT_LEN);
60        ffi::blake3_hash_many_avx512(
61            inputs.as_ptr() as *const *const u8,
62            inputs.len(),
63            N / BLOCK_LEN,
64            key.as_ptr(),
65            counter,
66            increment_counter.yes(),
67            flags,
68            flags_start,
69            flags_end,
70            out.as_mut_ptr(),
71        )
72    }
73}
74
75// Unsafe because this may only be called on platforms supporting AVX-512.
76#[cfg(unix)]
77pub unsafe fn xof_many(
78    cv: &CVWords,
79    block: &[u8; BLOCK_LEN],
80    block_len: u8,
81    counter: u64,
82    flags: u8,
83    out: &mut [u8],
84) {
85    unsafe {
86        debug_assert_eq!(0, out.len() % BLOCK_LEN, "whole blocks only");
87        ffi::blake3_xof_many_avx512(
88            cv.as_ptr(),
89            block.as_ptr(),
90            block_len,
91            counter,
92            flags,
93            out.as_mut_ptr(),
94            out.len() / BLOCK_LEN,
95        );
96    }
97}
98
99pub mod ffi {
100    extern "C" {
101        pub fn blake3_compress_in_place_avx512(
102            cv: *mut u32,
103            block: *const u8,
104            block_len: u8,
105            counter: u64,
106            flags: u8,
107        );
108        pub fn blake3_compress_xof_avx512(
109            cv: *const u32,
110            block: *const u8,
111            block_len: u8,
112            counter: u64,
113            flags: u8,
114            out: *mut u8,
115        );
116        pub fn blake3_hash_many_avx512(
117            inputs: *const *const u8,
118            num_inputs: usize,
119            blocks: usize,
120            key: *const u32,
121            counter: u64,
122            increment_counter: bool,
123            flags: u8,
124            flags_start: u8,
125            flags_end: u8,
126            out: *mut u8,
127        );
128        #[cfg(unix)]
129        pub fn blake3_xof_many_avx512(
130            cv: *const u32,
131            block: *const u8,
132            block_len: u8,
133            counter: u64,
134            flags: u8,
135            out: *mut u8,
136            outblocks: usize,
137        );
138    }
139}
140
141#[cfg(test)]
142mod test {
143    use super::*;
144
145    #[test]
146    fn test_compress() {
147        if !crate::platform::avx512_detected() {
148            return;
149        }
150        crate::test::test_compress_fn(compress_in_place, compress_xof);
151    }
152
153    #[test]
154    fn test_hash_many() {
155        if !crate::platform::avx512_detected() {
156            return;
157        }
158        crate::test::test_hash_many_fn(hash_many, hash_many);
159    }
160
161    #[cfg(unix)]
162    #[test]
163    fn test_xof_many() {
164        if !crate::platform::avx512_detected() {
165            return;
166        }
167        crate::test::test_xof_many_fn(xof_many);
168    }
169}
blake3/ffi_avx512.rs

blake3/
ffi_avx512.rs
