lance_linalg/
simd.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Poor-man's SIMD
5//!
6//! The difference between this implementation and [std::simd] is that
7//! this implementation holds the SIMD register directly, thus it exposes more
8//! optimization opportunity to use wide range of instructions available.
9//!
10//! Also, it gives us more control, for example, it is likely that we will have
11//! f16/bf16 support before the standard library does.
12//!
13//! The API are close to [std::simd] to make migration easier in the future.
14
15use std::ops::{Add, AddAssign, Mul, Sub, SubAssign};
16
17pub mod dist_table;
18pub mod f32;
19pub mod i32;
20pub mod u8;
21
22use num_traits::{Float, Num};
23use u8::u8x16;
24
25/// Lance SIMD lib
26///
27pub trait SIMD<T: Num + Copy, const N: usize>:
28    std::fmt::Debug
29    + AddAssign<Self>
30    + Add<Self, Output = Self>
31    + Mul<Self, Output = Self>
32    + Sub<Self, Output = Self>
33    + SubAssign<Self>
34    + Copy
35    + Clone
36    + Sized
37    + for<'a> From<&'a [T]>
38{
39    const LANES: usize = N;
40
41    /// Create a new instance with all lanes set to `val`.
42    fn splat(val: T) -> Self;
43
44    /// Create a new instance with all lanes set to zero.
45    fn zeros() -> Self;
46
47    /// Load aligned data from aligned memory.
48    ///
49    /// # Safety
50    ///
51    /// It crashes if the ptr is not aligned.
52    unsafe fn load(ptr: *const T) -> Self;
53
54    /// Load unaligned data from memory.
55    ///
56    /// # Safety
57    unsafe fn load_unaligned(ptr: *const T) -> Self;
58
59    /// Store the values to aligned memory.
60    ///
61    /// # Safety
62    ///
63    /// It crashes if the ptr is not aligned
64    unsafe fn store(&self, ptr: *mut T);
65
66    /// Store the values to unaligned memory.
67    ///
68    /// # Safety
69    unsafe fn store_unaligned(&self, ptr: *mut T);
70
71    /// Return the values as an array.
72    fn as_array(&self) -> [T; N] {
73        let mut arr = [T::zero(); N];
74        unsafe {
75            self.store_unaligned(arr.as_mut_ptr());
76        }
77        arr
78    }
79
80    /// Calculate the sum across this vector.
81    fn reduce_sum(&self) -> T;
82
83    /// Find the minimal value in the vector.
84    fn reduce_min(&self) -> T;
85
86    /// Return the minimal value of these two vectors.
87    fn min(&self, rhs: &Self) -> Self;
88
89    /// Find the index of value in the vector. If not found, return None.
90    fn find(&self, val: T) -> Option<i32>;
91}
92
93pub trait FloatSimd<F: Float, const N: usize>: SIMD<F, N> {
94    /// fused multiply-add
95    ///
96    /// c = a * b + c
97    fn multiply_add(&mut self, a: Self, b: Self);
98}
99
100pub trait Shuffle {
101    fn shuffle(&self, indices: u8x16) -> Self;
102}
lance_linalg/simd.rs

lance_linalg/
simd.rs
