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