OXYGENHUB
# c8d9059e4e70 oxygen / vendor / crc-fast / src / arch / aarch64 / aes.rs
7.9 KiB · Rust · last changed Strata: the cloud stores send through ureq, so every release target links 1 month ago source blame history raw
1 
// Copyright 2025 Don MacAskill. Licensed under MIT or Apache-2.0.
2
3 
//! This module provides AArch64-specific implementations of the ArchOps trait for architectures
4 
//! with AES support.
5
6 
#![cfg(target_arch = "aarch64")]
7
8 
use crate::traits::ArchOps;
9 
use core::arch::aarch64::*;
10
11 
// Tier-specific ArchOps implementations for AArch64
12
13 
/// Base AArch64 implementation with AES+NEON optimizations
14 
/// NEON is implicit with AES support on AArch64
15 
#[derive(Debug, Copy, Clone)]
16 
pub struct Aarch64AesOps;
17
18 
// Base implementation for AArch64 AES tier
19 
impl ArchOps for Aarch64AesOps {
20 
    type Vector = uint8x16_t;
21
22 
    #[inline]
23 
    #[target_feature(enable = "neon")]
24 
    unsafe fn create_vector_from_u64_pair(
25 
        &self,
26 
        high: u64,
27 
        low: u64,
28 
        reflected: bool,
29 
    ) -> Self::Vector {
30 
        // Note: AArch64 switches the order for reflected mode
31 
        if reflected {
32 
            self.load_key_pair(high, low)
33 
        } else {
34 
            self.load_key_pair(low, high)
35 
        }
36 
    }
37
38 
    #[inline]
39 
    #[target_feature(enable = "neon")]
40 
    unsafe fn create_vector_from_u64_pair_non_reflected(
41 
        &self,
42 
        high: u64,
43 
        low: u64,
44 
    ) -> Self::Vector {
45 
        self.load_key_pair(low, high)
46 
    }
47
48 
    #[inline]
49 
    #[target_feature(enable = "neon")]
50 
    unsafe fn create_vector_from_u64(&self, value: u64, high: bool) -> Self::Vector {
51 
        let mut result = vdupq_n_u64(0);
52 
        if high {
53 
            result = vsetq_lane_u64(value, result, 1); // Set high 64 bits
54 
        } else {
55 
            result = vsetq_lane_u64(value, result, 0); // Set low 64 bits
56 
        }
57
58 
        vreinterpretq_u8_u64(result)
59 
    }
60
61 
    #[inline]
62 
    #[target_feature(enable = "neon")]
63 
    unsafe fn extract_u64s(&self, vector: Self::Vector) -> [u64; 2] {
64 
        let x = vreinterpretq_u64_u8(vector);
65 
        [vgetq_lane_u64(x, 0), vgetq_lane_u64(x, 1)]
66 
    }
67
68 
    #[inline]
69 
    #[target_feature(enable = "neon")]
70 
    unsafe fn extract_poly64s(&self, vector: Self::Vector) -> [u64; 2] {
71 
        let x = vreinterpretq_p64_u8(vector);
72 
        [vgetq_lane_p64(x, 0), vgetq_lane_p64(x, 1)]
73 
    }
74
75 
    #[inline]
76 
    #[target_feature(enable = "neon")]
77 
    unsafe fn xor_vectors(&self, a: Self::Vector, b: Self::Vector) -> Self::Vector {
78 
        veorq_u8(a, b)
79 
    }
80
81 
    #[inline]
82 
    #[target_feature(enable = "neon")]
83 
    unsafe fn load_bytes(&self, ptr: *const u8) -> Self::Vector {
84 
        vld1q_u8(ptr)
85 
    }
86
87 
    #[inline]
88 
    #[target_feature(enable = "neon")]
89 
    unsafe fn load_aligned(&self, ptr: *const [u64; 2]) -> Self::Vector {
90 
        vreinterpretq_u8_u64(vld1q_u64(ptr as *const u64))
91 
    }
92
93 
    #[inline]
94 
    #[target_feature(enable = "neon")]
95 
    unsafe fn shuffle_bytes(&self, data: Self::Vector, mask: Self::Vector) -> Self::Vector {
96 
        vqtbl1q_u8(data, mask)
97 
    }
98
99 
    #[inline]
100 
    #[target_feature(enable = "neon")]
101 
    unsafe fn blend_vectors(
102 
        &self,
103 
        a: Self::Vector,
104 
        b: Self::Vector,
105 
        mask: Self::Vector,
106 
    ) -> Self::Vector {
107 
        let msb_mask = vcltq_s8(vreinterpretq_s8_u8(mask), vdupq_n_s8(0));
108 
        vbslq_u8(msb_mask, b, a)
109 
    }
110
111 
    #[inline]
112 
    #[target_feature(enable = "neon")]
113 
    unsafe fn shift_left_8(&self, vector: Self::Vector) -> Self::Vector {
114 
        vextq_u8(vdupq_n_u8(0), vector, 8)
115 
    }
116
117 
    #[inline]
118 
    #[target_feature(enable = "neon")]
119 
    unsafe fn set_all_bytes(&self, value: u8) -> Self::Vector {
120 
        vdupq_n_u8(value)
121 
    }
122
123 
    #[inline]
124 
    #[target_feature(enable = "neon")]
125 
    unsafe fn create_compare_mask(&self, vector: Self::Vector) -> Self::Vector {
126 
        vcltq_s8(vreinterpretq_s8_u8(vector), vdupq_n_s8(0))
127 
    }
128
129 
    #[inline]
130 
    #[target_feature(enable = "neon")]
131 
    unsafe fn and_vectors(&self, a: Self::Vector, b: Self::Vector) -> Self::Vector {
132 
        vandq_u8(a, b)
133 
    }
134
135 
    #[inline]
136 
    #[target_feature(enable = "neon")]
137 
    unsafe fn shift_right_32(&self, vector: Self::Vector) -> Self::Vector {
138 
        vextq_u8(vector, vdupq_n_u8(0), 4)
139 
    }
140
141 
    #[inline]
142 
    #[target_feature(enable = "neon")]
143 
    unsafe fn shift_left_32(&self, vector: Self::Vector) -> Self::Vector {
144 
        vextq_u8(vdupq_n_u8(0), vector, 12) // 16-4=12
145 
    }
146
147 
    #[inline]
148 
    #[target_feature(enable = "neon")]
149 
    unsafe fn create_vector_from_u32(&self, value: u32, high: bool) -> Self::Vector {
150 
        let mut result = vdupq_n_u64(0);
151 
        if high {
152 
            // For high=true, place in the high 32 bits of the high 64 bits
153 
            result = vreinterpretq_u64_u32(vsetq_lane_u32(value, vreinterpretq_u32_u64(result), 3));
154 
        } else {
155 
            // For high=false, place in the low 32 bits of the low 64 bits
156 
            result = vreinterpretq_u64_u32(vsetq_lane_u32(value, vreinterpretq_u32_u64(result), 0));
157 
        }
158 
        vreinterpretq_u8_u64(result)
159 
    }
160
161 
    #[inline]
162 
    #[target_feature(enable = "neon")]
163 
    unsafe fn shift_left_4(&self, vector: Self::Vector) -> Self::Vector {
164 
        vextq_u8(vdupq_n_u8(0), vector, 12) // 16-4=12
165 
    }
166
167 
    #[inline]
168 
    #[target_feature(enable = "neon")]
169 
    unsafe fn shift_right_4(&self, vector: Self::Vector) -> Self::Vector {
170 
        vextq_u8(vector, vdupq_n_u8(0), 4)
171 
    }
172
173 
    #[inline]
174 
    #[target_feature(enable = "neon")]
175 
    unsafe fn shift_right_8(&self, vector: Self::Vector) -> Self::Vector {
176 
        vextq_u8(vector, vdupq_n_u8(0), 8)
177 
    }
178
179 
    #[inline]
180 
    #[target_feature(enable = "neon")]
181 
    unsafe fn shift_right_5(&self, vector: Self::Vector) -> Self::Vector {
182 
        vextq_u8(vector, vdupq_n_u8(0), 5)
183 
    }
184
185 
    #[inline]
186 
    #[target_feature(enable = "neon")]
187 
    unsafe fn shift_right_6(&self, vector: Self::Vector) -> Self::Vector {
188 
        vextq_u8(vector, vdupq_n_u8(0), 6)
189 
    }
190
191 
    #[inline]
192 
    #[target_feature(enable = "neon")]
193 
    unsafe fn shift_right_7(&self, vector: Self::Vector) -> Self::Vector {
194 
        vextq_u8(vector, vdupq_n_u8(0), 7)
195 
    }
196
197 
    #[inline]
198 
    #[target_feature(enable = "neon")]
199 
    unsafe fn shift_right_12(&self, vector: Self::Vector) -> Self::Vector {
200 
        vextq_u8(vector, vdupq_n_u8(0), 12)
201 
    }
202
203 
    #[inline]
204 
    #[target_feature(enable = "neon")]
205 
    unsafe fn shift_left_12(&self, vector: Self::Vector) -> Self::Vector {
206 
        let low_32 = vgetq_lane_u32(vreinterpretq_u32_u8(vector), 0);
207 
        let result = vsetq_lane_u32(low_32, vdupq_n_u32(0), 3);
208 
        vreinterpretq_u8_u32(result)
209 
    }
210
211 
    #[inline]
212 
    #[target_feature(enable = "aes")]
213 
    unsafe fn carryless_mul_00(&self, a: Self::Vector, b: Self::Vector) -> Self::Vector {
214 
        vreinterpretq_u8_p128(vmull_p64(
215 
            vgetq_lane_p64(vreinterpretq_p64_u8(a), 0),
216 
            vgetq_lane_p64(vreinterpretq_p64_u8(b), 0),
217 
        ))
218 
    }
219
220 
    #[inline]
221 
    #[target_feature(enable = "aes")]
222 
    unsafe fn carryless_mul_01(&self, a: Self::Vector, b: Self::Vector) -> Self::Vector {
223 
        let a_low = vgetq_lane_p64(vreinterpretq_p64_u8(a), 1);
224 
        let b_high = vgetq_lane_p64(vreinterpretq_p64_u8(b), 0);
225 
        vreinterpretq_u8_p128(vmull_p64(a_low, b_high))
226 
    }
227
228 
    #[inline]
229 
    #[target_feature(enable = "aes")]
230 
    unsafe fn carryless_mul_10(&self, a: Self::Vector, b: Self::Vector) -> Self::Vector {
231 
        vreinterpretq_u8_p128(vmull_p64(
232 
            vgetq_lane_p64(vreinterpretq_p64_u8(a), 0),
233 
            vgetq_lane_p64(vreinterpretq_p64_u8(b), 1),
234 
        ))
235 
    }
236
237 
    #[inline]
238 
    #[target_feature(enable = "aes")]
239 
    unsafe fn carryless_mul_11(&self, a: Self::Vector, b: Self::Vector) -> Self::Vector {
240 
        vreinterpretq_u8_p128(vmull_p64(
241 
            vgetq_lane_p64(vreinterpretq_p64_u8(a), 1),
242 
            vgetq_lane_p64(vreinterpretq_p64_u8(b), 1),
243 
        ))
244 
    }
245
246 
    #[inline]
247 
    #[target_feature(enable = "neon")]
248 
    unsafe fn xor3_vectors(
249 
        &self,
250 
        a: Self::Vector,
251 
        b: Self::Vector,
252 
        c: Self::Vector,
253 
    ) -> Self::Vector {
254 
        veorq_u8(veorq_u8(a, b), c)
255 
    }
256 
}
257
258 
impl Aarch64AesOps {
259 
    #[inline]
260 
    #[target_feature(enable = "neon")]
261 
    unsafe fn load_key_pair(&self, idx1: u64, idx2: u64) -> uint8x16_t {
262 
        vreinterpretq_u8_u64(vld1q_u64([idx1, idx2].as_ptr()))
263 
    }
264 
}