Skip to content

fix(linalg): make SIMD slice conversions baseline-safe #8874

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lance-gatefixer wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(linalg): make SIMD slice conversions baseline-safe #8874

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 41 additions & 5 deletions rust/lance-linalg/src/simd/f32.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,10 @@ impl SIMD<f32, 8> for f32x8 {
unsafe fn load_unaligned(ptr: *const f32) -> Self {
#[cfg(target_arch = "x86_64")]
unsafe {
Self(_mm256_loadu_ps(ptr))
// Treat the register as plain data so this load remains valid on
// pre-AVX hosts. LLVM still lowers this to an AVX move when AVX is
// enabled for the caller.
Self(ptr.cast::<__m256>().read_unaligned())
}
#[cfg(target_arch = "aarch64")]
{
Expand Down Expand Up @@ -612,7 +615,11 @@ impl SIMD<f32, 16> for f32x16 {
unsafe fn load_unaligned(ptr: *const f32) -> Self {
#[cfg(target_arch = "x86_64")]
unsafe {
Self(_mm256_loadu_ps(ptr), _mm256_loadu_ps(ptr.add(8)))
// Treat the registers as plain data so this load remains valid on
// pre-AVX hosts. LLVM still lowers these to AVX moves when AVX is
// enabled for the caller.
let ptr = ptr.cast::<__m256>();
Self(ptr.read_unaligned(), ptr.add(1).read_unaligned())
}
#[cfg(target_arch = "aarch64")]
{
Expand Down Expand Up @@ -951,6 +958,32 @@ mod tests {
assert!(std::panic::catch_unwind(|| f32x16::from(&[0.0; 15][..])).is_err());
}

#[test]
fn test_slice_conversion_uses_cpu_baseline() {
let values8 = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
let values16 = [
1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0,
16.0,
];
let vectors = (
f32x8::from(&values8[..]),
f32x8::from(&values8),
f32x16::from(&values16[..]),
f32x16::from(&values16),
);

#[cfg(target_arch = "x86_64")]
if !std::is_x86_feature_detected!("avx") {
std::hint::black_box(vectors);
return;
}

assert_eq!(vectors.0.as_array(), values8);
assert_eq!(vectors.1.as_array(), values8);
assert_eq!(vectors.2.as_array(), values16);
assert_eq!(vectors.3.as_array(), values16);
}

#[test]
fn test_basic_ops() {
// Load / store / arithmetic on `f32x8` lower to AVX intrinsics, and
Expand Down Expand Up @@ -1124,12 +1157,15 @@ mod tests {
#[cfg(target_arch = "x86_64")]
#[test]
fn test_gather_scalar_x86() {
if !std::is_x86_feature_detected!("avx") {
return;
}
let a = (0..256).map(|f| f as f32).collect::<Vec<_>>();
let idx = [0_i32, 4, 8, 12, 16, 20, 24, 29];
let v = gather_scalar_x86(&a, &idx);

if !std::is_x86_feature_detected!("avx") {
std::hint::black_box(v);
return;
}

let expected = idx.map(|i| a[i as usize]);
assert_eq!(v.as_array(), expected);
}
Expand Down
34 changes: 32 additions & 2 deletions rust/lance-linalg/src/simd/f64.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,10 @@ impl SIMD<f64, 4> for f64x4 {
unsafe fn load_unaligned(ptr: *const f64) -> Self {
#[cfg(target_arch = "x86_64")]
unsafe {
Self(_mm256_loadu_pd(ptr))
// Treat the register as plain data so this load remains valid on
// pre-AVX hosts. LLVM still lowers this to an AVX move when AVX is
// enabled for the caller.
Self(ptr.cast::<__m256d>().read_unaligned())
}
#[cfg(target_arch = "aarch64")]
{
Expand Down Expand Up @@ -467,7 +470,11 @@ impl SIMD<f64, 8> for f64x8 {
unsafe fn load_unaligned(ptr: *const f64) -> Self {
#[cfg(target_arch = "x86_64")]
unsafe {
Self(_mm256_loadu_pd(ptr), _mm256_loadu_pd(ptr.add(4)))
// Treat the registers as plain data so this load remains valid on
// pre-AVX hosts. LLVM still lowers these to AVX moves when AVX is
// enabled for the caller.
let ptr = ptr.cast::<__m256d>();
Self(ptr.read_unaligned(), ptr.add(1).read_unaligned())
}
#[cfg(target_arch = "aarch64")]
unsafe {
Expand Down Expand Up @@ -750,6 +757,29 @@ mod tests {
assert!(std::panic::catch_unwind(|| f64x8::from(&[0.0; 7][..])).is_err());
}

#[test]
fn test_slice_conversion_uses_cpu_baseline() {
let values4 = [1.0_f64, 2.0, 3.0, 4.0];
let values8 = [1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
let vectors = (
f64x4::from(&values4[..]),
f64x4::from(&values4),
f64x8::from(&values8[..]),
f64x8::from(&values8),
);

#[cfg(target_arch = "x86_64")]
if !std::is_x86_feature_detected!("avx") {
std::hint::black_box(vectors);
return;
}

assert_eq!(vectors.0.as_array(), values4);
assert_eq!(vectors.1.as_array(), values4);
assert_eq!(vectors.2.as_array(), values8);
assert_eq!(vectors.3.as_array(), values8);
}

#[test]
fn test_f64x4_basic_ops() {
// The `f64x4` constructor / load / store / arithmetic paths all lower
Expand Down
21 changes: 20 additions & 1 deletion rust/lance-linalg/src/simd/i32.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,10 @@ impl SIMD<i32, 8> for i32x8 {
unsafe fn load_unaligned(ptr: *const i32) -> Self {
#[cfg(target_arch = "x86_64")]
unsafe {
Self(_mm256_loadu_si256(ptr as *const __m256i))
// Treat the register as plain data so this load remains valid on
// pre-AVX hosts. LLVM still lowers this to an AVX move when AVX is
// enabled for the caller.
Self(ptr.cast::<__m256i>().read_unaligned())
}
#[cfg(target_arch = "aarch64")]
{
Expand Down Expand Up @@ -353,6 +356,22 @@ mod tests {
assert!(std::panic::catch_unwind(|| i32x8::from(&[0; 7][..])).is_err());
}

#[test]
fn test_slice_conversion_uses_cpu_baseline() {
let values = [1_i32, 2, 3, 4, 5, 6, 7, 8];
let from_slice = i32x8::from(&values[..]);
let from_array = i32x8::from(&values);

#[cfg(target_arch = "x86_64")]
if !std::is_x86_feature_detected!("avx") {
std::hint::black_box((from_slice, from_array));
return;
}

assert_eq!(from_slice.as_array(), values);
assert_eq!(from_array.as_array(), values);
}

/// Lane-wise, low-32-bits multiplication is what all three arms promise, so
/// this runs everywhere: only the x86 feature check is arch-gated, matching
/// `f32.rs`'s and `f64.rs`'s test modules.
Expand Down
Loading