Skip to content

Commit

Permalink
Add a example benchmark for zipped iteration.
Browse files Browse the repository at this point in the history
  • Loading branch information
adamreichold committed Mar 17, 2022
1 parent d9443f1 commit b28f255
Showing 1 changed file with 85 additions and 0 deletions.
85 changes: 85 additions & 0 deletions benches/iter.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
#![feature(test)]
#![allow(deprecated)]

extern crate test;
use test::{black_box, Bencher};

use ndarray::Zip;
use numpy::{npyiter::NpyMultiIterBuilder, PyArray};
use pyo3::Python;

fn numpy_iter(bencher: &mut Bencher, size: usize) {
Python::with_gil(|py| {
let x = PyArray::<f64, _>::zeros(py, size, false);
let y = PyArray::<f64, _>::zeros(py, size, false);
let z = PyArray::<f64, _>::zeros(py, size, false);

let x = x.readonly();
let y = y.readonly();
let mut z = z.readwrite();

bencher.iter(|| {
let iter = NpyMultiIterBuilder::new()
.add_readonly(black_box(&x))
.add_readonly(black_box(&y))
.add_readwrite(black_box(&mut z))
.build()
.unwrap();

for (x, y, z) in iter {
*z = x + y;
}
});
});
}

#[bench]
fn numpy_iter_small(bencher: &mut Bencher) {
numpy_iter(bencher, 2_usize.pow(5));
}

#[bench]
fn numpy_iter_medium(bencher: &mut Bencher) {
numpy_iter(bencher, 2_usize.pow(10));
}

#[bench]
fn numpy_iter_large(bencher: &mut Bencher) {
numpy_iter(bencher, 2_usize.pow(15));
}

fn ndarray_iter(bencher: &mut Bencher, size: usize) {
Python::with_gil(|py| {
let x = PyArray::<f64, _>::zeros(py, size, false);
let y = PyArray::<f64, _>::zeros(py, size, false);
let z = PyArray::<f64, _>::zeros(py, size, false);

let x = x.readonly();
let y = y.readonly();
let mut z = z.readwrite();

bencher.iter(|| {
Zip::from(black_box(x.as_array()))
.and(black_box(y.as_array()))
.and(black_box(z.as_array_mut()))
.for_each(|x, y, z| {
*z = x + y;
});
});
});
}

#[bench]
fn ndarray_iter_small(bencher: &mut Bencher) {
ndarray_iter(bencher, 2_usize.pow(5));
}

#[bench]
fn ndarray_iter_medium(bencher: &mut Bencher) {
ndarray_iter(bencher, 2_usize.pow(10));
}

#[bench]
fn ndarray_iter_large(bencher: &mut Bencher) {
ndarray_iter(bencher, 2_usize.pow(15));
}

0 comments on commit b28f255

Please sign in to comment.