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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
//! The [`CaseSet`] API below is a safe and simplified version of the `case_set*` macros in `ctx.h`.
//!
//! The `case_set*` macros themselves replaced `memset`s in order to further optimize them
//! (in e3b5d4d044506f9e0e95e79b3de42fd94386cc61,
//! which performed the following optimizations:
//! 1. larger, inlinable writes for small power of 2 lengths
//! 2. aligned writes for the above
//! 3. 8-byte aligned fields [`BlockContext`]
//! * allows 8-byte aligned writes
//! * better cache boundary alignment
//!
//! (1) is easy to preserve in Rust, but (2) is difficult to do so without overhead,
//! as unaligned writes are UB, and so we'd need to check at runtime if they're aligned
//! (a runtime-determined `off`set is used, so we can't reasonably ensure this at compile-time).
//!
//! To more thoroughly check this, I ran the same benchmarks done in
//! e3b5d4d044506f9e0e95e79b3de42fd94386cc61, which introduced the `case_set*` macros:
//!
//! ```sh
//! cargo build --release && hyperfine './target/release/dav1d -i ./tests/large/chimera_8b_1080p.ivf -l 1000 -o /dev/null'
//! ```
//!
//! for 3 implementations:
//! 1. the original `case_set*` macros translated directly to `unsafe` Rust `fn`s
//! 2. the safe [`CaseSet`] implementation below using [`small_memset`] with its small powers of 2 optimization
//! 3. a safe [`CaseSet`] implementation using [`slice::fill`]/`memset` only
//!
//! The [`small_memset`] version was ~1.27% faster than the `case_set*` one,
//! and ~3.26% faster than the `memset` one.
//! The `case_set*` macros were also faster than `memset` in C by a similar margin,
//! meaning the `memset` option is the slowest in both C and Rust,
//! and since it was replaced with `case_set*` in C, we shouldn't use it in Rust.
//! Thus, the [`small_memset`] implementation seems optimal, as it:
//! * is the fastest of the Rust implementations
//! * is completely safe
//! * employs the same small powers of 2 optimization the `case_set*` implementation did
//! * is far simpler than the `case_set*` implementation, consisting of a `match` and array writes
//!
//! [`BlockContext`]: crate::src::env::BlockContext
use crate::src::disjoint_mut::AsMutPtr;
use crate::src::disjoint_mut::DisjointMut;
use std::iter::zip;
/// Perform a `memset` optimized for lengths that are small powers of 2.
///
/// For power of 2 lengths `<= UP_TO`,
/// the `memset` is done as an array write of that exactly (compile-time known) length.
/// If the length is not a power of 2 or `> UP_TO`,
/// then the `memset` is done by [`slice::fill`] (a `memset` call) if `WITH_DEFAULT` is `true`,
/// or else skipped if `WITH_DEFAULT` is `false`.
///
/// This optimizes for the common cases where `buf.len()` is a small power of 2,
/// where the array write is optimized as few and large stores as possible.
#[inline]
pub fn small_memset<T: Clone + Copy, const UP_TO: usize, const WITH_DEFAULT: bool>(
buf: &mut [T],
val: T,
) {
fn as_array<T: Clone + Copy, const N: usize>(buf: &mut [T]) -> &mut [T; N] {
buf.try_into().unwrap()
}
match buf.len() {
01 if UP_TO >= 01 => *as_array(buf) = [val; 01],
02 if UP_TO >= 02 => *as_array(buf) = [val; 02],
04 if UP_TO >= 04 => *as_array(buf) = [val; 04],
08 if UP_TO >= 08 => *as_array(buf) = [val; 08],
16 if UP_TO >= 16 => *as_array(buf) = [val; 16],
32 if UP_TO >= 32 => *as_array(buf) = [val; 32],
_ => {
if WITH_DEFAULT {
buf.fill(val)
}
}
}
}
pub struct CaseSetter<const UP_TO: usize, const WITH_DEFAULT: bool> {
offset: usize,
len: usize,
}
impl<const UP_TO: usize, const WITH_DEFAULT: bool> CaseSetter<UP_TO, WITH_DEFAULT> {
#[inline]
pub fn set<T: Clone + Copy>(&self, buf: &mut [T], val: T) {
small_memset::<T, UP_TO, WITH_DEFAULT>(&mut buf[self.offset..][..self.len], val);
}
/// # Safety
///
/// Caller must ensure that no elements of the written range are concurrently
/// borrowed (immutably or mutably) at all during the call to `set_disjoint`.
#[inline]
pub fn set_disjoint<T, V>(&self, buf: &DisjointMut<T>, val: V)
where
T: AsMutPtr<Target = V>,
V: Clone + Copy,
{
let mut buf = buf.index_mut(self.offset..self.offset + self.len);
small_memset::<V, UP_TO, WITH_DEFAULT>(&mut *buf, val);
}
}
/// The entrypoint to the [`CaseSet`] API.
///
/// `UP_TO` and `WITH_DEFAULT` are made const generic parameters rather than have multiple `case_set*` `fn`s,
/// and these are put in a separate `struct` so that these 2 generic parameters
/// can be manually specified while the ones on the methods are inferred.
pub struct CaseSet<const UP_TO: usize, const WITH_DEFAULT: bool>;
impl<const UP_TO: usize, const WITH_DEFAULT: bool> CaseSet<UP_TO, WITH_DEFAULT> {
/// Perform one case set.
///
/// This API is generic over the element type (`T`) rather than hardcoding `u8`,
/// as sometimes other types are used, though only `i8` is used currently.
///
/// The `len` and `offset` are supplied here and
/// applied to each `buf` passed to [`CaseSetter::set`] in `set_ctx`.
#[inline]
pub fn one<T, F>(ctx: T, len: usize, offset: usize, mut set_ctx: F)
where
F: FnMut(&CaseSetter<UP_TO, WITH_DEFAULT>, T),
{
set_ctx(&CaseSetter { offset, len }, ctx);
}
/// Perform many case sets in one call.
///
/// This allows specifying the `set_ctx` closure inline easily,
/// and also allows you to group the same args together.
///
/// The `lens`, `offsets`, and `dirs` are zipped and passed to [`CaseSet::one`],
/// where `dirs` can be an array of any type and whose elements are passed back to the `set_ctx` closure.
#[inline]
pub fn many<T, F, const N: usize>(
dirs: [T; N],
lens: [usize; N],
offsets: [usize; N],
mut set_ctx: F,
) where
F: FnMut(&CaseSetter<UP_TO, WITH_DEFAULT>, T),
{
for (dir, (len, offset)) in zip(dirs, zip(lens, offsets)) {
Self::one(dir, len, offset, &mut set_ctx);
}
}
}