Skip to content

Commit

Permalink
Improve Rust benchmarks (#115)
Browse files Browse the repository at this point in the history 
* playground with rayon

* improve benches

* not expose chunks_exact
  • Loading branch information
@edgarriba
edgarriba authored Aug 31, 2024
1 parent be59200 commit 1f8013f
Show file tree
Hide file tree
Showing 10 changed files with 334 additions and 142 deletions.
2 changes: 0 additions & 2 deletions crates/kornia-core/src/storage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,8 +125,6 @@ where
self.data.as_ptr()
}

// TODO: remove this method once we don't need it anymore because of
// ndarray::ArrayViewMut in kornia-imgproc
/// Returns the data pointer as a mutable pointer.
#[inline]
pub fn as_mut_ptr(&mut self) -> *mut T {
Expand Down
8 changes: 4 additions & 4 deletions crates/kornia-image/src/image.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -298,22 +298,22 @@ where

/// Get the number of columns of the image.
pub fn cols(&self) -> usize {
self.width()
self.shape[1]
}

/// Get the number of rows of the image.
pub fn rows(&self) -> usize {
self.height()
self.shape[0]
}

/// Get the width of the image in pixels.
pub fn width(&self) -> usize {
self.shape[1]
self.cols()
}

/// Get the height of the image in pixels.
pub fn height(&self) -> usize {
self.shape[0]
self.rows()
}

/// Get the number of channels in the image.
Expand Down
1 change: 1 addition & 0 deletions crates/kornia-imgproc/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ kornia.workspace = true
criterion = "0.5"
image = "0.25.1"
ndarray = { version = "0.15", features = ["rayon"] }
rayon = "1.10"

[[bench]]
name = "bench_color"
Expand Down
192 changes: 166 additions & 26 deletions crates/kornia-imgproc/benches/bench_color.rs
View file
Original file line number Diff line number Diff line change
@@ -1,22 +1,106 @@
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};

use kornia::imgproc::color::gray_from_rgb;
use kornia_image::{Image, ImageSize};
use kornia_image::Image;
use rayon::prelude::*;

// vanilla version
fn gray_iter(image: &Image<f32, 3>) -> Image<u8, 1> {
let data = vec![0u8; image.size().width * image.size().height];
let gray_image = Image::new(image.size(), data).unwrap();
for y in 0..image.height() {
for x in 0..image.width() {
let r = image.get_unchecked([y, x, 0]);
let g = image.get_unchecked([y, x, 1]);
let b = image.get_unchecked([y, x, 2]);
let _gray_pixel = (76. * r + 150. * g + 29. * b) / 255.;
// TODO: implement set_unchecked
fn gray_vanilla_get_unchecked(
src: &Image<f32, 3>,
dst: &mut Image<f32, 1>,
) -> Result<(), Box<dyn std::error::Error>> {
let data = dst.as_slice_mut();
let (cols, _rows) = (src.cols(), src.rows());

for y in 0..src.rows() {
for x in 0..src.cols() {
let r = src.get_unchecked([y, x, 0]);
let g = src.get_unchecked([y, x, 1]);
let b = src.get_unchecked([y, x, 2]);
let gray_pixel = (76. * r + 150. * g + 29. * b) / 255.;
data[y * cols + x] = gray_pixel;
}
}
gray_image
Ok(())
}

fn gray_slice_chunks_pixels(
src: &Image<f32, 3>,
dst: &mut Image<f32, 1>,
) -> Result<(), Box<dyn std::error::Error>> {
src.as_slice()
.chunks_exact(3)
.zip(dst.storage.as_mut_slice().chunks_exact_mut(1))
.for_each(|(src_chunk, dst_chunk)| {
let r = src_chunk[0];
let g = src_chunk[1];
let b = src_chunk[2];
dst_chunk[0] = (76. * r + 150. * g + 29. * b) / 255.;
});

Ok(())
}

fn gray_slice_chunks_pixels_parallel(
src: &Image<f32, 3>,
dst: &mut Image<f32, 1>,
) -> Result<(), Box<dyn std::error::Error>> {
src.as_slice()
.par_chunks_exact(3)
.zip(dst.as_slice_mut().par_chunks_exact_mut(1))
.for_each(|(src_chunk, dst_chunk)| {
let r = src_chunk[0];
let g = src_chunk[1];
let b = src_chunk[2];
dst_chunk[0] = (76. * r + 150. * g + 29. * b) / 255.;
});

Ok(())
}

fn gray_slice_chunks_rows(
src: &Image<f32, 3>,
dst: &mut Image<f32, 1>,
) -> Result<(), Box<dyn std::error::Error>> {
let num_cols = src.cols();
src.as_slice()
.chunks_exact(3 * num_cols)
.zip(dst.storage.as_mut_slice().chunks_exact_mut(num_cols))
.for_each(|(src_chunk, dst_chunk)| {
src_chunk
.chunks_exact(3)
.zip(dst_chunk.chunks_exact_mut(1))
.for_each(|(src_pixel, dst_pixel)| {
let r = src_pixel[0];
let g = src_pixel[1];
let b = src_pixel[2];
dst_pixel[0] = (76. * r + 150. * g + 29. * b) / 255.;
});
});

Ok(())
}
fn gray_slice_chunks_rows_parallel(
src: &Image<f32, 3>,
dst: &mut Image<f32, 1>,
) -> Result<(), Box<dyn std::error::Error>> {
let width = src.width();
src.as_slice()
.par_chunks_exact(3 * width)
.zip(dst.as_slice_mut().par_chunks_exact_mut(width))
.for_each(|(src_chunk, dst_chunk)| {
src_chunk
.chunks_exact(3)
.zip(dst_chunk.chunks_exact_mut(1))
.for_each(|(src_pixel, dst_pixel)| {
let r = src_pixel[0];
let g = src_pixel[1];
let b = src_pixel[2];
dst_pixel[0] = (76. * r + 150. * g + 29. * b) / 255.;
});
});

Ok(())
}

fn gray_image_crate(image: &Image<u8, 3>) -> Image<u8, 1> {
Expand All @@ -36,25 +120,81 @@ fn gray_image_crate(image: &Image<u8, 3>) -> Image<u8, 1> {

fn bench_grayscale(c: &mut Criterion) {
let mut group = c.benchmark_group("Grayscale");
let image_sizes = vec![(256, 224), (512, 448), (1024, 896)];

for (width, height) in image_sizes {
let id = format!("{}x{}", width, height);
for (width, height) in [(256, 224), (512, 448), (1024, 896)].iter() {
group.throughput(criterion::Throughput::Elements((*width * *height) as u64));

let parameter_string = format!("{}x{}", width, height);

// input image
let image_data = vec![0u8; width * height * 3];
let image = Image::new(ImageSize { width, height }, image_data).unwrap();
let image_size = [*width, *height].into();

let image = Image::new(image_size, image_data).unwrap();
let image_f32 = image.clone().cast::<f32>().unwrap();

// output image
let mut gray = Image::from_size_val(image.size(), 0.0).unwrap();
group.bench_with_input(BenchmarkId::new("zip", &id), &image_f32, |b, _i| {
b.iter(|| gray_from_rgb(black_box(&image_f32), black_box(&mut gray)))
});
group.bench_with_input(BenchmarkId::new("iter", &id), &image_f32, |b, i| {
b.iter(|| gray_iter(black_box(&i.clone())))
});
group.bench_with_input(BenchmarkId::new("image_crate", &id), &image, |b, i| {
b.iter(|| gray_image_crate(black_box(&i.clone())))
});
let gray = Image::from_size_val(image.size(), 0.0).unwrap();

group.bench_with_input(
BenchmarkId::new("vanilla_unchecked", &parameter_string),
&(&image_f32, &gray),
|b, i| {
let (src, mut dst) = (i.0, i.1.clone());
b.iter(|| black_box(gray_vanilla_get_unchecked(src, &mut dst)))
},
);

group.bench_with_input(
BenchmarkId::new("image_crate", &parameter_string),
&image,
|b, i| b.iter(|| black_box(gray_image_crate(i))),
);

group.bench_with_input(
BenchmarkId::new("ndarray_zip_par", &parameter_string),
&(&image_f32, &gray),
|b, i| {
let (src, mut dst) = (i.0, i.1.clone());
b.iter(|| black_box(gray_from_rgb(src, &mut dst)))
},
);

group.bench_with_input(
BenchmarkId::new("slice_chunks_pixels", &parameter_string),
&(&image_f32, &gray),
|b, i| {
let (src, mut dst) = (i.0, i.1.clone());
b.iter(|| black_box(gray_slice_chunks_pixels(src, &mut dst)))
},
);

group.bench_with_input(
BenchmarkId::new("slice_chunks_pixels_parallel", &parameter_string),
&(&image_f32, &gray),
|b, i| {
let (src, mut dst) = (i.0, i.1.clone());
b.iter(|| black_box(gray_slice_chunks_pixels_parallel(src, &mut dst)))
},
);

group.bench_with_input(
BenchmarkId::new("slice_chunks_rows", &parameter_string),
&(&image_f32, &gray),
|b, i| {
let (src, mut dst) = (i.0, i.1.clone());
b.iter(|| black_box(gray_slice_chunks_rows(src, &mut dst)))
},
);

group.bench_with_input(
BenchmarkId::new("slice_chunks_rows_parallel", &parameter_string),
&(&image_f32, &gray),
|b, i| {
let (src, mut dst) = (i.0, i.1.clone());
b.iter(|| black_box(gray_slice_chunks_rows_parallel(src, &mut dst)))
},
);
}
group.finish();
}
Expand Down
25 changes: 14 additions & 11 deletions crates/kornia-imgproc/benches/bench_metrics.rs
View file
Original file line number Diff line number Diff line change
@@ -1,22 +1,25 @@
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};

use kornia::{
image::{Image, ImageSize},
imgproc::metrics,
};
use kornia::{image::Image, imgproc::metrics};

fn bench_mse(c: &mut Criterion) {
let mut group = c.benchmark_group("mse");
let image_sizes = vec![(256, 224), (512, 448), (1024, 896)];

for (width, height) in image_sizes {
let image_size = ImageSize { width, height };
let id = format!("{}x{}", width, height);
for (width, height) in [(256, 224), (512, 448), (1024, 896)].iter() {
group.throughput(criterion::Throughput::Elements((*width * *height) as u64));

let parameter_string = format!("{}x{}", width, height);

// input image
let image_size = [*width, *height].into();
let image = Image::<u8, 3>::new(image_size, vec![0u8; width * height * 3]).unwrap();
let image_f32 = image.cast::<f32>().unwrap();
group.bench_with_input(BenchmarkId::new("mapv", &id), &image_f32, |b, i| {
b.iter(|| metrics::mse(black_box(i), black_box(i)))
});

group.bench_with_input(
BenchmarkId::new("mse_map", &parameter_string),
&image_f32,
|b, i| b.iter(|| black_box(metrics::mse(i, i))),
);
}
group.finish();
}
Expand Down
74 changes: 46 additions & 28 deletions crates/kornia-imgproc/benches/bench_resize.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,44 +23,62 @@ fn resize_image_crate(image: Image<u8, 3>, new_size: ImageSize) -> Image<u8, 3>
}

fn bench_resize(c: &mut Criterion) {
let mut group = c.benchmark_group("resize");
let image_sizes = vec![(256, 224), (512, 448), (1024, 896)];
let mut group = c.benchmark_group("Resize");

for (width, height) in [(256, 224), (512, 448), (1024, 896)].iter() {
group.throughput(criterion::Throughput::Elements((*width * *height) as u64));

let parameter_string = format!("{}x{}", width, height);

for (width, height) in image_sizes {
let image_size = ImageSize { width, height };
let id = format!("{}x{}", width, height);
// input image
let image_size = [*width, *height].into();
let image = Image::<u8, 3>::new(image_size, vec![0u8; width * height * 3]).unwrap();
let image_f32 = image.clone().cast::<f32>().unwrap();

// output image
let new_size = ImageSize {
width: width / 2,
height: height / 2,
};
let mut out_f32 = Image::<f32, 3>::from_size_val(new_size, 0.0).unwrap();
let mut out_u8 = Image::<u8, 3>::from_size_val(new_size, 0).unwrap();

group.bench_with_input(BenchmarkId::new("native", &id), &image_f32, |b, i| {
b.iter(|| {
resize::resize_native(
black_box(i),
black_box(&mut out_f32),
InterpolationMode::Nearest,
)
})
});
group.bench_with_input(BenchmarkId::new("image_rs", &id), &image, |b, i| {
b.iter(|| resize_image_crate(black_box(i.clone()), new_size))
});
group.bench_with_input(BenchmarkId::new("fast", &id), &image, |b, i| {
b.iter(|| {
resize::resize_fast(
black_box(i),
black_box(&mut out_u8),
InterpolationMode::Nearest,
)
})
});
let out_f32 = Image::<f32, 3>::from_size_val(new_size, 0.0).unwrap();
let out_u8 = Image::<u8, 3>::from_size_val(new_size, 0).unwrap();

group.bench_with_input(
BenchmarkId::new("image_rs", &parameter_string),
&image,
|b, i| b.iter(|| resize_image_crate(black_box(i.clone()), new_size)),
);

group.bench_with_input(
BenchmarkId::new("ndarray_zip_par", &parameter_string),
&(&image_f32, &out_f32),
|b, i| {
let (src, mut dst) = (i.0, i.1.clone());
b.iter(|| {
black_box(resize::resize_native(
src,
&mut dst,
InterpolationMode::Nearest,
))
})
},
);

group.bench_with_input(
BenchmarkId::new("fast_resize_lib", &parameter_string),
&(image, out_u8),
|b, i| {
let (src, mut dst) = (i.0.clone(), i.1.clone());
b.iter(|| {
black_box(resize::resize_fast(
&src,
&mut dst,
InterpolationMode::Nearest,
))
})
},
);
}
group.finish();
}
Expand Down
Loading

0 comments on commit 1f8013f

Please sign in to comment.