Skip to content

Commit 833896d

Browse files
adamreicholdadamreichold
authored
Merge pull request #292 from PyO3/copy-from-vec23
Benchmark and optimize PyArray2/3::from_vec2/3
2 parents 3324bc7 + 0b3a83e commit 833896d

File tree

7 files changed

+201
-68
lines changed

7 files changed

+201
-68
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Changelog
22

33
- Unreleased
4+
- Deprecate `PyArray::from_exact_iter` after optimizing `PyArray::from_iter`. ([#292](https://github.com/PyO3/rust-numpy/pull/292))
45

56
- v0.16.2
67
- Fix build on platforms where `c_char` is `u8` like Linux/AArch64. ([#296](https://github.com/PyO3/rust-numpy/pull/296))

‎benches/array.rs

Lines changed: 52 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use test::{black_box, Bencher};
55

66
use std::ops::Range;
77

8-
use numpy::PyArray1;
8+
use numpy::{PyArray1, PyArray2, PyArray3};
99
use pyo3::{Python, ToPyObject};
1010

1111
struct Iter(Range<usize>);
@@ -65,6 +65,7 @@ fn from_exact_iter(bencher: &mut Bencher, size: usize) {
6565
iter_with_gil(bencher, |py| {
6666
let iter = black_box(ExactIter(0..size));
6767

68+
#[allow(deprecated)]
6869
PyArray1::from_exact_iter(py, iter);
6970
});
7071
}
@@ -134,6 +135,56 @@ fn from_object_slice_large(bencher: &mut Bencher) {
134135
from_object_slice(bencher, 2_usize.pow(15));
135136
}
136137

138+
fn from_vec2(bencher: &mut Bencher, size: usize) {
139+
let vec2 = vec![vec![0; size]; size];
140+
141+
iter_with_gil(bencher, |py| {
142+
let vec2 = black_box(&vec2);
143+
144+
PyArray2::from_vec2(py, vec2).unwrap();
145+
});
146+
}
147+
148+
#[bench]
149+
fn from_vec2_small(bencher: &mut Bencher) {
150+
from_vec2(bencher, 2_usize.pow(3));
151+
}
152+
153+
#[bench]
154+
fn from_vec2_medium(bencher: &mut Bencher) {
155+
from_vec2(bencher, 2_usize.pow(5));
156+
}
157+
158+
#[bench]
159+
fn from_vec2_large(bencher: &mut Bencher) {
160+
from_vec2(bencher, 2_usize.pow(8));
161+
}
162+
163+
fn from_vec3(bencher: &mut Bencher, size: usize) {
164+
let vec3 = vec![vec![vec![0; size]; size]; size];
165+
166+
iter_with_gil(bencher, |py| {
167+
let vec3 = black_box(&vec3);
168+
169+
PyArray3::from_vec3(py, vec3).unwrap();
170+
});
171+
}
172+
173+
#[bench]
174+
fn from_vec3_small(bencher: &mut Bencher) {
175+
from_vec3(bencher, 2_usize.pow(2));
176+
}
177+
178+
#[bench]
179+
fn from_vec3_medium(bencher: &mut Bencher) {
180+
from_vec3(bencher, 2_usize.pow(4));
181+
}
182+
183+
#[bench]
184+
fn from_vec3_large(bencher: &mut Bencher) {
185+
from_vec3(bencher, 2_usize.pow(5));
186+
}
187+
137188
fn iter_with_gil(bencher: &mut Bencher, mut f: impl FnMut(Python)) {
138189
Python::with_gil(|py| {
139190
bencher.iter(|| {

‎src/array.rs

Lines changed: 65 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ use pyo3::{
1919
Python, ToPyObject,
2020
};
2121

22+
use crate::cold;
2223
use crate::convert::{ArrayExt, IntoPyArray, NpyIndex, ToNpyDims, ToPyArray};
2324
use crate::dtype::{Element, PyArrayDescr};
2425
use crate::error::{DimensionalityError, FromVecError, NotContiguousError, TypeError};
@@ -954,14 +955,8 @@ impl<T: Element> PyArray<T, Ix1> {
954955
pub fn from_slice<'py>(py: Python<'py>, slice: &[T]) -> &'py Self {
955956
unsafe {
956957
let array = PyArray::new(py, [slice.len()], false);
957-
if T::IS_COPY {
958-
ptr::copy_nonoverlapping(slice.as_ptr(), array.data(), slice.len());
959-
} else {
960-
let data_ptr = array.data();
961-
for (i, item) in slice.iter().enumerate() {
962-
data_ptr.add(i).write(item.clone());
963-
}
964-
}
958+
let mut data_ptr = array.data();
959+
clone_elements(slice, &mut data_ptr);
965960
array
966961
}
967962
}
@@ -978,6 +973,7 @@ impl<T: Element> PyArray<T, Ix1> {
978973
/// assert_eq!(pyarray.readonly().as_slice().unwrap(), &[1, 2, 3, 4, 5]);
979974
/// });
980975
/// ```
976+
#[inline(always)]
981977
pub fn from_vec<'py>(py: Python<'py>, vec: Vec<T>) -> &'py Self {
982978
vec.into_pyarray(py)
983979
}
@@ -988,34 +984,45 @@ impl<T: Element> PyArray<T, Ix1> {
988984
/// # Example
989985
/// ```
990986
/// use numpy::PyArray;
991-
/// use std::collections::BTreeSet;
992-
/// let vec = vec![1, 2, 3, 4, 5];
993-
/// pyo3::Python::with_gil(|py| {
994-
/// let pyarray = PyArray::from_exact_iter(py, vec.iter().map(|&x| x));
987+
/// use pyo3::Python;
988+
///
989+
/// Python::with_gil(|py| {
990+
/// let pyarray = PyArray::from_exact_iter(py, [1, 2, 3, 4, 5].into_iter().copied());
995991
/// assert_eq!(pyarray.readonly().as_slice().unwrap(), &[1, 2, 3, 4, 5]);
996992
/// });
997993
/// ```
998-
pub fn from_exact_iter(py: Python<'_>, iter: impl ExactSizeIterator<Item = T>) -> &Self {
999-
let data = iter.collect::<Box<[_]>>();
1000-
data.into_pyarray(py)
994+
#[deprecated(
995+
note = "`from_exact_iter` is deprecated as it does not provide any benefit over `from_iter`."
996+
)]
997+
#[inline(always)]
998+
pub fn from_exact_iter<I>(py: Python<'_>, iter: I) -> &Self
999+
where
1000+
I: IntoIterator<Item = T>,
1001+
I::IntoIter: ExactSizeIterator,
1002+
{
1003+
Self::from_iter(py, iter)
10011004
}
10021005

1003-
/// Construct one-dimension PyArray from a type which implements
1004-
/// [`IntoIterator`](https://doc.rust-lang.org/std/iter/trait.IntoIterator.html).
1006+
/// Construct one-dimension PyArray from a type which implements [`IntoIterator`].
10051007
///
1006-
/// If no reliable [`size_hint`](https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.size_hint) is available,
1008+
/// If no reliable [`size_hint`][Iterator::size_hint] is available,
10071009
/// this method can allocate memory multiple time, which can hurt performance.
10081010
///
10091011
/// # Example
1012+
///
10101013
/// ```
10111014
/// use numpy::PyArray;
1012-
/// let set: std::collections::BTreeSet<u32> = [4, 3, 2, 5, 1].into_iter().cloned().collect();
1013-
/// pyo3::Python::with_gil(|py| {
1014-
/// let pyarray = PyArray::from_iter(py, set);
1015-
/// assert_eq!(pyarray.readonly().as_slice().unwrap(), &[1, 2, 3, 4, 5]);
1015+
/// use pyo3::Python;
1016+
///
1017+
/// Python::with_gil(|py| {
1018+
/// let pyarray = PyArray::from_iter(py, "abcde".chars().map(u32::from));
1019+
/// assert_eq!(pyarray.readonly().as_slice().unwrap(), &[97, 98, 99, 100, 101]);
10161020
/// });
10171021
/// ```
1018-
pub fn from_iter(py: Python<'_>, iter: impl IntoIterator<Item = T>) -> &Self {
1022+
pub fn from_iter<I>(py: Python<'_>, iter: I) -> &Self
1023+
where
1024+
I: IntoIterator<Item = T>,
1025+
{
10191026
let data = iter.into_iter().collect::<Vec<_>>();
10201027
data.into_pyarray(py)
10211028
}
@@ -1095,19 +1102,18 @@ impl<T: Element> PyArray<T, Ix2> {
10951102
/// });
10961103
/// ```
10971104
pub fn from_vec2<'py>(py: Python<'py>, v: &[Vec<T>]) -> Result<&'py Self, FromVecError> {
1098-
let last_len = v.last().map_or(0, |v| v.len());
1099-
for v in v {
1100-
if v.len() != last_len {
1101-
return Err(FromVecError::new(v.len(), last_len));
1102-
}
1103-
}
1104-
let dims = [v.len(), last_len];
1105+
let len2 = v.first().map_or(0, |v| v.len());
1106+
let dims = [v.len(), len2];
1107+
// SAFETY: The result of `Self::new` is always safe to drop.
11051108
unsafe {
11061109
let array = Self::new(py, dims, false);
1107-
for (y, vy) in v.iter().enumerate() {
1108-
for (x, vyx) in vy.iter().enumerate() {
1109-
array.uget_raw([y, x]).write(vyx.clone());
1110+
let mut data_ptr = array.data();
1111+
for v in v {
1112+
if v.len() != len2 {
1113+
cold();
1114+
return Err(FromVecError::new(v.len(), len2));
11101115
}
1116+
clone_elements(v, &mut data_ptr);
11111117
}
11121118
Ok(array)
11131119
}
@@ -1135,28 +1141,24 @@ impl<T: Element> PyArray<T, Ix3> {
11351141
/// });
11361142
/// ```
11371143
pub fn from_vec3<'py>(py: Python<'py>, v: &[Vec<Vec<T>>]) -> Result<&'py Self, FromVecError> {
1138-
let len2 = v.last().map_or(0, |v| v.len());
1139-
for v in v {
1140-
if v.len() != len2 {
1141-
return Err(FromVecError::new(v.len(), len2));
1142-
}
1143-
}
1144-
let len3 = v.last().map_or(0, |v| v.last().map_or(0, |v| v.len()));
1145-
for v in v {
1146-
for v in v {
1147-
if v.len() != len3 {
1148-
return Err(FromVecError::new(v.len(), len3));
1149-
}
1150-
}
1151-
}
1144+
let len2 = v.first().map_or(0, |v| v.len());
1145+
let len3 = v.first().map_or(0, |v| v.first().map_or(0, |v| v.len()));
11521146
let dims = [v.len(), len2, len3];
1147+
// SAFETY: The result of `Self::new` is always safe to drop.
11531148
unsafe {
11541149
let array = Self::new(py, dims, false);
1155-
for (z, vz) in v.iter().enumerate() {
1156-
for (y, vzy) in vz.iter().enumerate() {
1157-
for (x, vzyx) in vzy.iter().enumerate() {
1158-
array.uget_raw([z, y, x]).write(vzyx.clone());
1150+
let mut data_ptr = array.data();
1151+
for v in v {
1152+
if v.len() != len2 {
1153+
cold();
1154+
return Err(FromVecError::new(v.len(), len2));
1155+
}
1156+
for v in v {
1157+
if v.len() != len3 {
1158+
cold();
1159+
return Err(FromVecError::new(v.len(), len3));
11591160
}
1161+
clone_elements(v, &mut data_ptr);
11601162
}
11611163
}
11621164
Ok(array)
@@ -1298,6 +1300,18 @@ impl<T: Element + AsPrimitive<f64>> PyArray<T, Ix1> {
12981300
}
12991301
}
13001302

1303+
unsafe fn clone_elements<T: Element>(elems: &[T], data_ptr: &mut *mut T) {
1304+
if T::IS_COPY {
1305+
ptr::copy_nonoverlapping(elems.as_ptr(), *data_ptr, elems.len());
1306+
*data_ptr = data_ptr.add(elems.len());
1307+
} else {
1308+
for elem in elems {
1309+
data_ptr.write(elem.clone());
1310+
*data_ptr = data_ptr.add(1);
1311+
}
1312+
}
1313+
}
1314+
13011315
#[cfg(test)]
13021316
mod tests {
13031317
use super::*;

‎src/convert.rs

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -163,9 +163,10 @@ where
163163
);
164164
unsafe {
165165
let array = PyArray::<A, _>::new_(py, dim, strides.as_ptr(), 0);
166-
let data_ptr = array.data();
167-
for (i, item) in self.iter().enumerate() {
168-
data_ptr.add(i).write(item.clone());
166+
let mut data_ptr = array.data();
167+
for item in self.iter() {
168+
data_ptr.write(item.clone());
169+
data_ptr = data_ptr.add(1);
169170
}
170171
array
171172
}

‎src/lib.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,10 @@ mod sealed {
8282
pub trait Sealed {}
8383
}
8484

85+
#[cold]
86+
#[inline(always)]
87+
fn cold() {}
88+
8589
/// Create a [`PyArray`] with one, two or three dimensions.
8690
///
8791
/// This macro is backed by [`ndarray::array`].

0 commit comments

Comments
 (0)