1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
use std::mem::{MaybeUninit, transmute};
use crate::Elem;
use crate::ops::BitOps;
/// Utility for incrementally filling an uninitialized slice, one SIMD vector
/// at a time.
pub struct SliceWriter<'a, T> {
buf: &'a mut [MaybeUninit<T>],
n_init: usize,
}
impl<'a, T: Elem> SliceWriter<'a, T> {
/// Create a writer which initializes elements of `buf`.
pub fn new(buf: &'a mut [MaybeUninit<T>]) -> Self {
SliceWriter { buf, n_init: 0 }
}
/// Initialize the next `ops.len()` elements of the slice from the contents
/// of SIMD vector `xs`.
///
/// Panics if the slice does not have space for `ops.len()` elements.
pub fn write_vec<O: BitOps<T>>(&mut self, ops: O, xs: O::Simd) {
let written = ops.store_uninit(xs, &mut self.buf[self.n_init..]);
self.n_init += written.len();
}
/// Initialize the next `N * ops.len()` elements of the slice from the
/// contents of the SIMD vectors `xs`.
///
/// This is equivalent to calling [`write_vec`](Self::write_vec) for each
/// vector, but performs a single bounds check for the whole batch, which is
/// useful when writing several vectors per loop iteration.
///
/// Panics if the slice does not have space for `N * ops.len()` elements.
pub fn write_vecs<O: BitOps<T>, const N: usize>(&mut self, ops: O, xs: [O::Simd; N]) {
let written = ops.store_many_uninit(xs, &mut self.buf[self.n_init..]);
self.n_init += written.len();
}
/// Initialize the next element of the slice from `x`.
///
/// Panics if the slice does not have space for writing any more elements.
pub fn write_scalar(&mut self, x: T) {
self.buf[self.n_init].write(x);
self.n_init += 1;
}
/// Finish writing the slice and return the initialized portion.
pub fn into_mut_slice(self) -> &'a mut [T] {
let init = &mut self.buf[0..self.n_init];
// Safety: All elements in `init` have been initialized.
unsafe { transmute::<&mut [MaybeUninit<T>], &mut [T]>(init) }
}
}
#[cfg(test)]
mod tests {
use std::mem::MaybeUninit;
use crate::ops::BitOps;
use crate::{Isa, SimdOp, SliceWriter};
#[test]
fn test_slice_writer() {
struct MemCopy<'src, 'dest> {
src: &'src [f32],
dest: &'dest mut [MaybeUninit<f32>],
}
impl<'src, 'dest> SimdOp for MemCopy<'src, 'dest> {
type Output = &'dest mut [f32];
fn eval<I: Isa>(self, isa: I) -> &'dest mut [f32] {
let ops = isa.f32();
let mut src_chunks = self.src.chunks_exact(ops.len());
let mut dest_writer = SliceWriter::new(self.dest);
for chunk in src_chunks.by_ref() {
let xs = ops.load(chunk);
dest_writer.write_vec(ops, xs);
}
for x in src_chunks.remainder() {
dest_writer.write_scalar(*x);
}
dest_writer.into_mut_slice()
}
}
// Length which should cover the vectorized body and tail cases for
// every ISA.
let len = 17;
let src: Vec<_> = (0..len).map(|x| x as f32).collect();
let mut dest = Vec::with_capacity(src.len());
let copied = MemCopy {
src: &src,
dest: dest.spare_capacity_mut(),
}
.dispatch();
assert_eq!(copied, src);
}
}