Skip to content

Commit 4076ba8

Browse files
committed
Implement scatter/gather with new pointer vector and add tests
1 parent 7e96f5d commit 4076ba8

File tree

10 files changed

+277
-161
lines changed

10 files changed

+277
-161
lines changed

crates/core_simd/src/cast.rs

Lines changed: 108 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,41 @@
1-
use crate::simd::SimdElement;
1+
use crate::simd::{intrinsics, LaneCount, Simd, SimdElement, SupportedLaneCount};
22

33
/// Supporting trait for `Simd::cast`. Typically doesn't need to be used directly.
4-
pub trait SimdCast<Target: SimdElement>: SimdElement {}
4+
pub trait SimdCast<Target: SimdElement>: SimdElement {
5+
#[doc(hidden)]
6+
fn cast<const LANES: usize>(x: Simd<Self, LANES>) -> Simd<Target, LANES>
7+
where
8+
LaneCount<LANES>: SupportedLaneCount;
9+
}
510

611
macro_rules! into_number {
12+
{ $from:ty, $to:ty } => {
13+
impl SimdCast<$to> for $from {
14+
fn cast<const LANES: usize>(x: Simd<Self, LANES>) -> Simd<$to, LANES>
15+
where
16+
LaneCount<LANES>: SupportedLaneCount,
17+
{
18+
// Safety: simd_as can handle numeric conversions
19+
unsafe { intrinsics::simd_as(x) }
20+
}
21+
}
22+
};
723
{ $($type:ty),* } => {
824
$(
9-
impl SimdCast<i8> for $type {}
10-
impl SimdCast<i16> for $type {}
11-
impl SimdCast<i32> for $type {}
12-
impl SimdCast<i64> for $type {}
13-
impl SimdCast<isize> for $type {}
14-
15-
impl SimdCast<u8> for $type {}
16-
impl SimdCast<u16> for $type {}
17-
impl SimdCast<u32> for $type {}
18-
impl SimdCast<u64> for $type {}
19-
impl SimdCast<usize> for $type {}
20-
21-
impl SimdCast<f32> for $type {}
22-
impl SimdCast<f64> for $type {}
25+
into_number! { $type, i8 }
26+
into_number! { $type, i16 }
27+
into_number! { $type, i32 }
28+
into_number! { $type, i64 }
29+
into_number! { $type, isize }
30+
31+
into_number! { $type, u8 }
32+
into_number! { $type, u16 }
33+
into_number! { $type, u32 }
34+
into_number! { $type, u64 }
35+
into_number! { $type, usize }
36+
37+
into_number! { $type, f32 }
38+
into_number! { $type, f64 }
2339
)*
2440
}
2541
}
@@ -29,17 +45,85 @@ into_number! { i8, i16, i32, i64, isize, u8, u16, u32, u64, usize, f32, f64 }
2945
macro_rules! into_pointer {
3046
{ $($type:ty),* } => {
3147
$(
32-
impl<T> SimdCast<$type> for *const T {}
33-
impl<T> SimdCast<$type> for *mut T {}
34-
impl<T> SimdCast<*const T> for $type {}
35-
impl<T> SimdCast<*mut T> for $type {}
48+
impl<T> SimdCast<$type> for *const T {
49+
fn cast<const LANES: usize>(x: Simd<Self, LANES>) -> Simd<$type, LANES>
50+
where
51+
LaneCount<LANES>: SupportedLaneCount,
52+
{
53+
// Safety: transmuting isize to pointers is safe
54+
let x: Simd<isize, LANES> = unsafe { core::mem::transmute_copy(&x) };
55+
x.cast()
56+
}
57+
}
58+
impl<T> SimdCast<$type> for *mut T {
59+
fn cast<const LANES: usize>(x: Simd<Self, LANES>) -> Simd<$type, LANES>
60+
where
61+
LaneCount<LANES>: SupportedLaneCount,
62+
{
63+
// Safety: transmuting isize to pointers is safe
64+
let x: Simd<isize, LANES> = unsafe { core::mem::transmute_copy(&x) };
65+
x.cast()
66+
}
67+
}
68+
impl<T> SimdCast<*const T> for $type {
69+
fn cast<const LANES: usize>(x: Simd<$type, LANES>) -> Simd<*const T, LANES>
70+
where
71+
LaneCount<LANES>: SupportedLaneCount,
72+
{
73+
let x: Simd<isize, LANES> = x.cast();
74+
// Safety: transmuting isize to pointers is safe
75+
unsafe { core::mem::transmute_copy(&x) }
76+
}
77+
}
78+
impl<T> SimdCast<*mut T> for $type {
79+
fn cast<const LANES: usize>(x: Simd<$type, LANES>) -> Simd<*mut T, LANES>
80+
where
81+
LaneCount<LANES>: SupportedLaneCount,
82+
{
83+
let x: Simd<isize, LANES> = x.cast();
84+
// Safety: transmuting isize to pointers is safe
85+
unsafe { core::mem::transmute_copy(&x) }
86+
}
87+
}
3688
)*
3789
}
3890
}
3991

4092
into_pointer! { i8, i16, i32, i64, isize, u8, u16, u32, u64, usize }
4193

42-
impl<T, U> SimdCast<*const T> for *const U {}
43-
impl<T, U> SimdCast<*const T> for *mut U {}
44-
impl<T, U> SimdCast<*mut T> for *const U {}
45-
impl<T, U> SimdCast<*mut T> for *mut U {}
94+
impl<T, U> SimdCast<*const T> for *const U {
95+
fn cast<const LANES: usize>(x: Simd<*const U, LANES>) -> Simd<*const T, LANES>
96+
where
97+
LaneCount<LANES>: SupportedLaneCount,
98+
{
99+
// Safety: transmuting pointers is safe
100+
unsafe { core::mem::transmute_copy(&x) }
101+
}
102+
}
103+
impl<T, U> SimdCast<*const T> for *mut U {
104+
fn cast<const LANES: usize>(x: Simd<*mut U, LANES>) -> Simd<*const T, LANES>
105+
where
106+
LaneCount<LANES>: SupportedLaneCount,
107+
{
108+
// Safety: transmuting pointers is safe
109+
unsafe { core::mem::transmute_copy(&x) }
110+
}
111+
}
112+
impl<T, U> SimdCast<*mut T> for *const U {
113+
fn cast<const LANES: usize>(x: Simd<*const U, LANES>) -> Simd<*mut T, LANES>
114+
where
115+
LaneCount<LANES>: SupportedLaneCount,
116+
{
117+
// Safety: transmuting pointers is safe
118+
unsafe { core::mem::transmute_copy(&x) }
119+
}
120+
}
121+
impl<T, U> SimdCast<*mut T> for *mut U {
122+
fn cast<const LANES: usize>(x: Simd<*mut U, LANES>) -> Simd<*mut T, LANES>
123+
where
124+
LaneCount<LANES>: SupportedLaneCount,
125+
{
126+
// Safety: transmuting pointers is safe
127+
unsafe { core::mem::transmute_copy(&x) }
128+
}
129+
}

crates/core_simd/src/elements/const_ptr.rs

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ use crate::simd::{LaneCount, Mask, Simd, SimdPartialEq, SupportedLaneCount};
33

44
/// Operations on SIMD vectors of constant pointers.
55
pub trait SimdConstPtr: Copy + Sealed {
6-
/// Vector type representing the pointers as bits.
7-
type Bits;
6+
/// Vector of usize with the same number of lanes.
7+
type Usize;
88

99
/// Vector of mutable pointers to the same type.
1010
type MutPtr;
@@ -18,11 +18,15 @@ pub trait SimdConstPtr: Copy + Sealed {
1818
/// Changes constness without changing the type.
1919
fn as_mut(self) -> Self::MutPtr;
2020

21-
/// Cast pointers to raw bits.
22-
fn to_bits(self) -> Self::Bits;
21+
/// Gets the "address" portion of the pointer.
22+
///
23+
/// Equivalent to calling [`pointer::addr`] on each lane.
24+
fn addr(self) -> Self::Usize;
2325

24-
/// Cast raw bits to pointers.
25-
fn from_bits(bits: Self::Bits) -> Self;
26+
/// Calculates the offset from a pointer using wrapping arithmetic.
27+
///
28+
/// Equivalent to calling [`pointer::wrapping_add`] on each lane.
29+
fn wrapping_add(self, count: Self::Usize) -> Self;
2630
}
2731

2832
impl<T, const LANES: usize> Sealed for Simd<*const T, LANES> where
@@ -34,23 +38,29 @@ impl<T, const LANES: usize> SimdConstPtr for Simd<*const T, LANES>
3438
where
3539
LaneCount<LANES>: SupportedLaneCount,
3640
{
37-
type Bits = Simd<usize, LANES>;
41+
type Usize = Simd<usize, LANES>;
3842
type MutPtr = Simd<*mut T, LANES>;
3943
type Mask = Mask<isize, LANES>;
4044

45+
#[inline]
4146
fn is_null(self) -> Self::Mask {
4247
Simd::splat(core::ptr::null()).simd_eq(self)
4348
}
4449

50+
#[inline]
4551
fn as_mut(self) -> Self::MutPtr {
4652
self.cast()
4753
}
4854

49-
fn to_bits(self) -> Self::Bits {
55+
#[inline]
56+
fn addr(self) -> Self::Usize {
5057
self.cast()
5158
}
5259

53-
fn from_bits(bits: Self::Bits) -> Self {
54-
bits.cast()
60+
#[inline]
61+
fn wrapping_add(self, count: Self::Usize) -> Self {
62+
let addr = self.addr() + (count * Simd::splat(core::mem::size_of::<T>()));
63+
// Safety: transmuting usize to pointers is safe, even if accessing those pointers isn't.
64+
unsafe { core::mem::transmute_copy(&addr) }
5565
}
5666
}

crates/core_simd/src/elements/mut_ptr.rs

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ use crate::simd::{LaneCount, Mask, Simd, SimdPartialEq, SupportedLaneCount};
33

44
/// Operations on SIMD vectors of mutable pointers.
55
pub trait SimdMutPtr: Copy + Sealed {
6-
/// Vector type representing the pointers as bits.
7-
type Bits;
6+
/// Vector of usize with the same number of lanes.
7+
type Usize;
88

99
/// Vector of constant pointers to the same type.
1010
type ConstPtr;
@@ -18,11 +18,15 @@ pub trait SimdMutPtr: Copy + Sealed {
1818
/// Changes constness without changing the type.
1919
fn as_const(self) -> Self::ConstPtr;
2020

21-
/// Cast pointers to raw bits.
22-
fn to_bits(self) -> Self::Bits;
21+
/// Gets the "address" portion of the pointer.
22+
///
23+
/// Equivalent to calling [`pointer::addr`] on each lane.
24+
fn addr(self) -> Self::Usize;
2325

24-
/// Cast raw bits to pointers.
25-
fn from_bits(bits: Self::Bits) -> Self;
26+
/// Calculates the offset from a pointer using wrapping arithmetic.
27+
///
28+
/// Equivalent to calling [`pointer::wrapping_add`] on each lane.
29+
fn wrapping_add(self, count: Self::Usize) -> Self;
2630
}
2731

2832
impl<T, const LANES: usize> Sealed for Simd<*mut T, LANES> where LaneCount<LANES>: SupportedLaneCount
@@ -32,23 +36,29 @@ impl<T, const LANES: usize> SimdMutPtr for Simd<*mut T, LANES>
3236
where
3337
LaneCount<LANES>: SupportedLaneCount,
3438
{
35-
type Bits = Simd<usize, LANES>;
39+
type Usize = Simd<usize, LANES>;
3640
type ConstPtr = Simd<*const T, LANES>;
3741
type Mask = Mask<isize, LANES>;
3842

43+
#[inline]
3944
fn is_null(self) -> Self::Mask {
4045
Simd::splat(core::ptr::null_mut()).simd_eq(self)
4146
}
4247

48+
#[inline]
4349
fn as_const(self) -> Self::ConstPtr {
4450
self.cast()
4551
}
4652

47-
fn to_bits(self) -> Self::Bits {
53+
#[inline]
54+
fn addr(self) -> Self::Usize {
4855
self.cast()
4956
}
5057

51-
fn from_bits(bits: Self::Bits) -> Self {
52-
bits.cast()
58+
#[inline]
59+
fn wrapping_add(self, count: Self::Usize) -> Self {
60+
let addr = self.addr() + (count * Simd::splat(core::mem::size_of::<T>()));
61+
// Safety: transmuting usize to pointers is safe, even if accessing those pointers isn't.
62+
unsafe { core::mem::transmute_copy(&addr) }
5363
}
5464
}

crates/core_simd/src/eq.rs

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
use crate::simd::{intrinsics, LaneCount, Mask, Simd, SimdElement, SupportedLaneCount};
1+
use crate::simd::{
2+
intrinsics, LaneCount, Mask, Simd, SimdConstPtr, SimdElement, SimdMutPtr, SupportedLaneCount,
3+
};
24

35
/// Parallel `PartialEq`.
46
pub trait SimdPartialEq {
@@ -80,16 +82,12 @@ where
8082

8183
#[inline]
8284
fn simd_eq(self, other: Self) -> Self::Mask {
83-
// Safety: `self` is a vector, and the result of the comparison
84-
// is always a valid mask.
85-
unsafe { Mask::from_int_unchecked(intrinsics::simd_eq(self, other)) }
85+
self.addr().simd_eq(other.addr())
8686
}
8787

8888
#[inline]
8989
fn simd_ne(self, other: Self) -> Self::Mask {
90-
// Safety: `self` is a vector, and the result of the comparison
91-
// is always a valid mask.
92-
unsafe { Mask::from_int_unchecked(intrinsics::simd_ne(self, other)) }
90+
self.addr().simd_ne(other.addr())
9391
}
9492
}
9593

@@ -101,15 +99,11 @@ where
10199

102100
#[inline]
103101
fn simd_eq(self, other: Self) -> Self::Mask {
104-
// Safety: `self` is a vector, and the result of the comparison
105-
// is always a valid mask.
106-
unsafe { Mask::from_int_unchecked(intrinsics::simd_eq(self, other)) }
102+
self.addr().simd_eq(other.addr())
107103
}
108104

109105
#[inline]
110106
fn simd_ne(self, other: Self) -> Self::Mask {
111-
// Safety: `self` is a vector, and the result of the comparison
112-
// is always a valid mask.
113-
unsafe { Mask::from_int_unchecked(intrinsics::simd_ne(self, other)) }
107+
self.addr().simd_ne(other.addr())
114108
}
115109
}

0 commit comments

Comments
 (0)