1#![cfg_attr(not(feature = "std"), no_std)]
38#![deny(missing_docs)]
39#![allow(
40 unused_unsafe,
41 clippy::too_many_arguments,
42 clippy::needless_range_loop,
43 clippy::assign_op_pattern,
44 clippy::manual_memcpy
45)]
46
47extern crate alloc;
48
49pub use hermes_simd_core::{
50 align::{Aligned, Alignment, Unaligned},
51 arch::SimdArch,
52 bitboard::{BitBoardKernel, BitBoardView},
53 compute::{ComputeReduce, ComputeView},
55 cow::{ArchivedPacked4Cow, ArchivedSimdCow, Packed4CowResolver, SimdCow, SimdCowResolver},
56 current_numa_node,
57 execution::{ExecutionMode, Masked, Unmasked},
58 iter::{SimdChunks, SimdChunksMut, ZipChunks, ZipChunksMut},
60 kernel::SimdKernel,
61 mask::BitMask,
62 refresh_numa_node,
63 scalar::{CastFrom, CastTo, FloatElement, Scalar as SimdScalar},
64 vec::AlignedVec,
65 verify_numa_locality,
66 view::{Mask, SimdError, SimdView, TileMatrixMultiply, TileView, Vector},
67 Abs,
69 Add,
70 BitAnd,
71 BitOr,
72 BitXor,
73 Clamp,
74 Div,
75 Dot,
76 ElementOp,
77 Exclusive,
78 FmaAdd,
80 Inclusive,
81 Mul,
82 Neg,
83 NumaBinding,
84 Popcount,
85 Product,
86 RecipSqrt,
87 ReductionOp,
89 ScanAdd,
90 ScanMax,
91 ScanMin,
92 ScanMode,
93 ScanMul,
94 ScanOp,
96 Sqrt,
97 Sub,
98 Sum,
99 UnaryOp,
100};
101
102pub use hermes_simd_core::sparse::{
104 BlockedCoo,
105 BlockedCooData,
106 CowFormat,
108 Csr,
109 CsrData,
110 DenseWithMask,
111 DenseWithMaskData,
112 OwnedBlockedCoo,
113 OwnedCsr,
115 OwnedDenseWithMask,
116 OwnedSellP,
117 SellP,
118 SellPData,
119 SparseCow,
121 SparseFormat,
122 SparseOps,
123 SparseSpMv,
124 SparseView,
125 Validated,
126 ValidatedData,
127};
128
129pub use hermes_simd_core::tiling::{tiled_dot, tiled_gemv, TilingPolicy, TilingStrategy};
131
132pub use hermes_simd_core::tensor::{ColMajor, RowMajor, TensorCow, TensorError, TensorView};
134
135pub use hermes_simd_intrinsics::{
137 Avx2, Avx512, AvxVnni, HybridSwarMagic, Hyperbola, KoggeStone, Magic, Neon, Scalar, SveArch,
138 Swar, SwarUtils,
139};
140
141#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
142pub use hermes_simd_intrinsics::{
143 AmxBatchSession, AmxBf16, AmxConfig, AmxInt8, AmxSession, AmxSessionError,
144};
145
146pub use eunomia::{
147 unpack_bf4_to_bf16, unpack_bf4_to_bf16_packed, unpack_bf8_to_bf16, unpack_f4_to_f32,
148 unpack_f4_to_f32_packed, unpack_f8_to_f32, Bf16, Bf4, Bf8, Packable4, Packed4Cow, Packed4Iter,
149 Packed4Slice, Packed4SliceMut, Packed4Vec, PackedBf4Cow, PackedBf4Slice, PackedBf4SliceMut,
150 PackedBf4Vec, PackedF4Cow, PackedF4Slice, PackedF4SliceMut, PackedF4Vec, F16, F32, F4, F64, F8,
151 I16, I32, I8,
152};
153
154pub use hermes_simd_types::{
156 MaskBf16, MaskBf4, MaskBf8, MaskF16, MaskF32, MaskF4, MaskF64, MaskF8, MaskI16, MaskI32,
157 MaskI8, PreferredArch, ScalarBf16, ScalarBf4, ScalarBf8, ScalarF16, ScalarF32, ScalarF4,
158 ScalarF64, ScalarF8, ScalarI16, ScalarI32, ScalarI8, ScalarMaskF32, ScalarMaskF64, SimdBf16,
159 SimdBf4, SimdBf8, SimdF16, SimdF32, SimdF4, SimdF64, SimdF8, SimdI16, SimdI32, SimdI8,
160 SimdMaskBf16, SimdMaskBf4, SimdMaskBf8, SimdMaskF16, SimdMaskF32, SimdMaskF4, SimdMaskF64,
161 SimdMaskF8, SimdMaskI16, SimdMaskI32, SimdMaskI8, VectorBf16, VectorBf4, VectorBf8, VectorF16,
162 VectorF32, VectorF4, VectorF64, VectorF8, VectorI16, VectorI32, VectorI8,
163};
164
165#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
166pub use hermes_simd_types::{
167 Avx2Bf16, Avx2Bf4, Avx2Bf8, Avx2F16, Avx2F32, Avx2F4, Avx2F64, Avx2F8, Avx2I16, Avx2I32,
168 Avx2I8, Avx2MaskBf16, Avx2MaskF16, Avx2MaskF32, Avx2MaskF64, Avx512Bf16, Avx512Bf4, Avx512Bf8,
169 Avx512F16, Avx512F32, Avx512F4, Avx512F64, Avx512F8, Avx512I16, Avx512I32, Avx512I8,
170 Avx512MaskBf16, Avx512MaskF16, Avx512MaskF32, Avx512MaskF64,
171};
172
173#[cfg(target_arch = "aarch64")]
174pub use hermes_simd_types::{
175 NeonBf16, NeonBf4, NeonBf8, NeonF16, NeonF32, NeonF4, NeonF64, NeonF8, NeonI16, NeonI32,
176 NeonI8, NeonMaskBf16, NeonMaskF16, NeonMaskF32, NeonMaskF64,
177};
178
179pub mod attacks;
181pub mod cpu;
183pub mod dispatcher;
185
186pub mod tile_matmul;
188
189pub mod dispatch;
191pub mod target;
193
194pub use attacks::{bishop_attacks, queen_attacks, rook_attacks};
195pub use cpu::{has_fma3, AmxSupport, Avx512Support, FmaSupport};
196pub use dispatcher::{AdaptiveDispatcher, DispatchDecision};
197pub use target::{dispatch_view_mut_to, dispatch_view_to, TargetId};
198pub use tile_matmul::{
199 dispatch_tile_matmul, gemm, unpack_int4, widen_I8_to_I16, widen_I8_to_I32, widen_i8_to_i16,
200 widen_i8_to_i32, TiledGemm,
201};
202
203pub use dispatch::{
206 abs_max,
207 abs_sum,
208 argmax,
209 argmin,
210 axpy,
211 axpy_mul,
212 axpy_rows,
213 axpy_rows_batch,
214 dot,
215 elementwise_add,
216 elementwise_div,
217 elementwise_mul,
218 elementwise_sub,
219 gemv,
220 gemv_strided,
221 gemv_transpose,
222 gemv_transpose_strided,
223 interleaved_complex_dot,
224 interleaved_complex_dot_runtime,
225 interleaved_complex_mul_assign,
226 interleaved_complex_mul_assign_runtime,
227 masked_add,
228 masked_dot,
229 masked_sum,
230 max,
231 min,
232 ntt_butterfly_stage_u64,
233 reduce_popcount,
235 reduce_popcount_and,
236 reduce_popcount_or,
237 reduce_popcount_xor,
238 scale,
239 spmv_bcoo,
240 spmv_csr,
242 spmv_dense_masked,
243 spmv_sellp,
244 sum,
245 tiled_gemm,
246 SimdOps,
248};
249
250pub enum DispatchedView<'a, T, Align = Unaligned, Mode = Unmasked, Ref = &'a [T]>
252where
253 Align: hermes_simd_core::align::Alignment,
254 Mode: hermes_simd_core::execution::ExecutionMode,
255 Ref: core::ops::Deref<Target = [T]>,
256{
257 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
259 Avx512(SimdView<'a, T, Avx512, Align, Mode, Ref>),
260 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
262 Avx2(SimdView<'a, T, Avx2, Align, Mode, Ref>),
263 #[cfg(target_arch = "aarch64")]
265 Neon(SimdView<'a, T, Neon, Align, Mode, Ref>),
266 Scalar(SimdView<'a, T, Scalar, Align, Mode, Ref>),
268}
269
270#[inline]
272#[allow(unreachable_code)]
273pub fn dispatch_view<'a, T, Align>(
274 data: &'a [T],
275) -> Option<DispatchedView<'a, T, Align, Unmasked, &'a [T]>>
276where
277 T: FloatElement,
278 Align: Alignment,
279{
280 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
281 {
282 #[cfg(feature = "std")]
283 {
284 if std::is_x86_feature_detected!("avx512f") {
285 return SimdView::<T, Avx512, Align, Unmasked, &'a [T]>::new(data)
286 .map(DispatchedView::Avx512);
287 }
288 if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
289 return SimdView::<T, Avx2, Align, Unmasked, &'a [T]>::new(data)
290 .map(DispatchedView::Avx2);
291 }
292 }
293 #[cfg(not(feature = "std"))]
294 {
295 if cfg!(target_feature = "avx512f") {
296 return SimdView::<T, Avx512, Align, Unmasked, &'a [T]>::new(data)
297 .map(DispatchedView::Avx512);
298 }
299 if cfg!(target_feature = "avx2") && cfg!(target_feature = "fma") {
300 return SimdView::<T, Avx2, Align, Unmasked, &'a [T]>::new(data)
301 .map(DispatchedView::Avx2);
302 }
303 }
304 }
305 #[cfg(target_arch = "aarch64")]
306 {
307 return SimdView::<T, Neon, Align, Unmasked, &'a [T]>::new(data).map(DispatchedView::Neon);
308 }
309 SimdView::<T, Scalar, Align, Unmasked, &'a [T]>::new(data).map(DispatchedView::Scalar)
310}
311
312#[inline]
314#[allow(unreachable_code)]
315pub fn dispatch_view_mut<'a, T, Align>(
316 data: &'a mut [T],
317) -> Option<DispatchedView<'a, T, Align, Unmasked, &'a mut [T]>>
318where
319 T: FloatElement,
320 Align: Alignment,
321{
322 #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
323 {
324 #[cfg(feature = "std")]
325 {
326 if std::is_x86_feature_detected!("avx512f") {
327 return SimdView::<T, Avx512, Align, Unmasked, &'a mut [T]>::new_mut(data)
328 .map(DispatchedView::Avx512);
329 }
330 if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
331 return SimdView::<T, Avx2, Align, Unmasked, &'a mut [T]>::new_mut(data)
332 .map(DispatchedView::Avx2);
333 }
334 }
335 #[cfg(not(feature = "std"))]
336 {
337 if cfg!(target_feature = "avx512f") {
338 return SimdView::<T, Avx512, Align, Unmasked, &'a mut [T]>::new_mut(data)
339 .map(DispatchedView::Avx512);
340 }
341 if cfg!(target_feature = "avx2") && cfg!(target_feature = "fma") {
342 return SimdView::<T, Avx2, Align, Unmasked, &'a mut [T]>::new_mut(data)
343 .map(DispatchedView::Avx2);
344 }
345 }
346 }
347 #[cfg(target_arch = "aarch64")]
348 {
349 return SimdView::<T, Neon, Align, Unmasked, &'a mut [T]>::new_mut(data)
350 .map(DispatchedView::Neon);
351 }
352 SimdView::<T, Scalar, Align, Unmasked, &'a mut [T]>::new_mut(data).map(DispatchedView::Scalar)
353}
354
355pub trait SimdCowExt<T: SimdScalar, Arch: SimdArch + SimdKernel<T>, Align: Alignment> {
357 fn transform_vectors<F>(&mut self, f: F)
362 where
363 F: FnMut(Vector<T, Arch>) -> Vector<T, Arch>;
364}
365
366impl<'a, T, Arch, Align> SimdCowExt<T, Arch, Align> for SimdCow<'a, T, Arch, Align>
367where
368 T: SimdScalar,
369 Arch: SimdArch + SimdKernel<T>,
370 Align: Alignment,
371{
372 #[inline]
373 fn transform_vectors<F>(&mut self, mut f: F)
374 where
375 F: FnMut(Vector<T, Arch>) -> Vector<T, Arch>,
376 {
377 let owned_vec = self.to_mut();
378 let len = owned_vec.len();
379 let lane_count = Arch::LANE_COUNT;
380 let simd_len = (len / lane_count) * lane_count;
381 let slice = owned_vec.as_mut_slice();
382
383 let mut i = 0;
384 while i < simd_len {
385 unsafe {
386 let ptr = slice.as_mut_ptr().add(i);
387 let vec = if Align::IS_ALIGNED {
388 Vector::load_aligned(ptr)
389 } else {
390 Vector::load_unaligned(ptr)
391 };
392 let res = f(vec);
393 if Align::IS_ALIGNED {
394 res.store_aligned(ptr);
395 } else {
396 res.store_unaligned(ptr);
397 }
398 }
399 i += lane_count;
400 }
401
402 if simd_len < len {
403 debug_assert!(
407 lane_count <= 64,
408 "tail_buf[64] too small: lane_count={lane_count}"
409 );
410 let mut tail_buf = [core::mem::MaybeUninit::<T>::uninit(); 64];
411 let tail_len = len - simd_len;
412 unsafe {
413 for idx in 0..tail_len {
414 tail_buf[idx].write(slice[simd_len + idx]);
415 }
416 for idx in tail_len..lane_count {
417 tail_buf[idx].write(T::ZERO);
418 }
419 let vec = Vector::load_unaligned(tail_buf.as_ptr() as *const T);
420 let res = f(vec);
421 res.store_unaligned(tail_buf.as_mut_ptr() as *mut T);
422 for idx in 0..tail_len {
423 slice[simd_len + idx] = tail_buf[idx].assume_init();
424 }
425 }
426 }
427 }
428}
429
430pub trait Packed4CowExt<'a, T: Packable4> {
432 fn unpack_to_cow<Arch, Align>(&self) -> SimdCow<'static, T::Unpacked, Arch, Align>
434 where
435 Arch: SimdArch,
436 Align: Alignment;
437}
438
439impl<'a, T: Packable4> Packed4CowExt<'a, T> for Packed4Cow<'a, T> {
440 #[inline]
441 fn unpack_to_cow<Arch, Align>(&self) -> SimdCow<'static, T::Unpacked, Arch, Align>
442 where
443 Arch: SimdArch,
444 Align: Alignment,
445 {
446 let len = self.len();
447 let mut dest = AlignedVec::with_capacity(len);
448 unsafe {
449 dest.set_len(len);
450 }
451 let view = self.as_view();
452 let n = view.len().min(dest.len());
453 let even_len = (n / 2) * 2;
454 T::unpack_slice_packed(
455 &view.as_packed_slice()[..even_len / 2],
456 &mut dest[..even_len],
457 );
458 if n % 2 != 0 {
459 if let Some(b) = view.get(n - 1) {
460 dest[n - 1] = T::unpack_single(b);
461 }
462 }
463 SimdCow::Owned(dest)
464 }
465}