use criterion to benchmark cava

juno_samples/cava/benches/cava_bench.rs

+#![feature(concat_idents)]
+
 use criterion::{criterion_group, criterion_main, Criterion};
 
-use juno_cava::{cava_harness, CavaInputs};
+use hercules_rt::{runner, HerculesImmBoxTo};
+
+use juno_cava::*;
+
+juno_build::juno!("cava");
+
+fn cava_bench(c: &mut Criterion) {
+    let mut group = c.benchmark_group("cava bench");
+    group.sample_size(10);
+
+    let args = CavaInputs {
+        input: "examples/raw_tulip-small.bin".to_string(),
+        output: None,
+        verify: false,
+        output_verify: None,
+        cam_model: "cam_models/NikonD7000".to_string(),
+        crop_rows: Some(144),
+        crop_cols: Some(192),
+    };
+    let (raw_image, cam_model) = make_raw_image_and_cam_model(&args);
+    let (rows, cols, num_ctrl_pts, image, tstw, ctrl_pts, weights, coefs, tonemap) =
+        prepare_hercules_inputs(&raw_image, &cam_model);
+    let mut r = runner!(cava);
+
+    group.bench_function("cava bench small", |b| {
+        b.iter(|| {
+            async_std::task::block_on(r.run(
+                rows as u64,
+                cols as u64,
+                num_ctrl_pts as u64,
+                image.to(),
+                tstw.to(),
+                ctrl_pts.to(),
+                weights.to(),
+                coefs.to(),
+                tonemap.to(),
+            ));
+        })
+    });
+
+    let args = CavaInputs {
+        input: "examples/raw_tulips.bin".to_string(),
+        output: None,
+        verify: true,
+        output_verify: None,
+        cam_model: "cam_models/NikonD7000".to_string(),
+        crop_rows: None,
+        crop_cols: None,
+    };
+    let (raw_image, cam_model) = make_raw_image_and_cam_model(&args);
+    let (rows, cols, num_ctrl_pts, image, tstw, ctrl_pts, weights, coefs, tonemap) =
+        prepare_hercules_inputs(&raw_image, &cam_model);
+    let mut r = runner!(cava);
 
-fn cava_bench_small(c: &mut Criterion) {
-    c.bench_function("cava_bench_small", |b| {
+    group.bench_function("cava bench full", |b| {
         b.iter(|| {
-            cava_harness(CavaInputs {
-                input: "examples/raw_tulip-small.bin".to_string(),
-                output: None,
-                verify: false,
-                output_verify: None,
-                cam_model: "cam_models/NikonD7000".to_string(),
-                crop_rows: Some(144),
-                crop_cols: Some(192),
-            })
+            async_std::task::block_on(r.run(
+                rows as u64,
+                cols as u64,
+                num_ctrl_pts as u64,
+                image.to(),
+                tstw.to(),
+                ctrl_pts.to(),
+                weights.to(),
+                coefs.to(),
+                tonemap.to(),
+            ));
         })
     });
 }
 
-criterion_group!(benches, cava_bench_small);
+criterion_group!(benches, cava_bench);
 criterion_main!(benches);

juno_samples/cava/src/lib.rs

@@ -4,9 +4,9 @@ mod camera_model;
 mod cava_rust;
 mod image_proc;
 
-use self::camera_model::*;
-use self::cava_rust::CHAN;
-use self::image_proc::*;
+pub use self::camera_model::*;
+pub use self::cava_rust::CHAN;
+pub use self::image_proc::*;
 
 use hercules_rt::{runner, HerculesImmBox, HerculesImmBoxTo, HerculesMutBox};
 
@@ -16,82 +16,59 @@ use clap::Parser;
 
 juno_build::juno!("cava");
 
-// Individual lifetimes are not needed in this example but should probably be generated for
-// flexibility
-async fn safe_run<'a, 'b: 'a, 'c: 'a, 'd: 'a, 'e: 'a, 'f: 'a, 'g: 'a>(
-    runner: &'a mut HerculesRunner_cava,
-    r: u64,
-    c: u64,
-    num_ctrl_pts: u64,
-    input: &'b HerculesImmBox<'b, u8>,
-    tstw: &'c HerculesImmBox<'c, f32>,
-    ctrl_pts: &'d HerculesImmBox<'d, f32>,
-    weights: &'e HerculesImmBox<'e, f32>,
-    coefs: &'f HerculesImmBox<'f, f32>,
-    tonemap: &'g HerculesImmBox<'g, f32>,
-) -> HerculesMutBox<'a, u8> {
-    HerculesMutBox::from(
-        runner
-            .run(
-                r,
-                c,
-                num_ctrl_pts,
-                input.to(),
-                tstw.to(),
-                ctrl_pts.to(),
-                weights.to(),
-                coefs.to(),
-                tonemap.to(),
-            )
-            .await,
-    )
+pub fn make_raw_image_and_cam_model(args: &CavaInputs) -> (RawImage, CamModel) {
+    let raw_image =
+        read_raw(&args.input, args.crop_rows, args.crop_cols).expect("Error loading image");
+    let cam_model = load_cam_model(&args.cam_model, CHAN).expect("Error loading camera model");
+    println!(
+        "Running cava with {} rows, {} columns, and {} control points.",
+        raw_image.rows, raw_image.cols, cam_model.num_ctrl_pts
+    );
+    (raw_image, cam_model)
 }
 
-fn run_cava(
-    rows: usize,
-    cols: usize,
-    num_ctrl_pts: usize,
-    image: &[u8],
-    tstw: &[f32],
-    ctrl_pts: &[f32],
-    weights: &[f32],
-    coefs: &[f32],
-    tonemap: &[f32],
-) -> Box<[u8]> {
-    assert_eq!(image.len(), CHAN * rows * cols);
-    assert_eq!(tstw.len(), CHAN * CHAN);
-    assert_eq!(ctrl_pts.len(), num_ctrl_pts * CHAN);
-    assert_eq!(weights.len(), num_ctrl_pts * CHAN);
-    assert_eq!(coefs.len(), 4 * CHAN);
-    assert_eq!(tonemap.len(), 256 * CHAN);
-
-    let image = HerculesImmBox::from(image);
-    let tstw = HerculesImmBox::from(tstw);
-    let ctrl_pts = HerculesImmBox::from(ctrl_pts);
-    let weights = HerculesImmBox::from(weights);
-    let coefs = HerculesImmBox::from(coefs);
-    let tonemap = HerculesImmBox::from(tonemap);
-
-    let mut r = runner!(cava);
-
-    async_std::task::block_on(async {
-        safe_run(
-            &mut r,
-            rows as u64,
-            cols as u64,
-            num_ctrl_pts as u64,
-            &image,
-            &tstw,
-            &ctrl_pts,
-            &weights,
-            &coefs,
-            &tonemap,
-        )
-        .await
-    })
-    .as_slice()
-    .to_vec()
-    .into_boxed_slice()
+pub fn prepare_hercules_inputs<'a, 'b>(
+    raw_image: &'a RawImage,
+    cam_model: &'b CamModel,
+) -> (
+    usize,
+    usize,
+    usize,
+    HerculesImmBox<'a, u8>,
+    HerculesImmBox<'b, f32>,
+    HerculesImmBox<'b, f32>,
+    HerculesImmBox<'b, f32>,
+    HerculesImmBox<'b, f32>,
+    HerculesImmBox<'b, f32>,
+) {
+    assert_eq!(
+        raw_image.pixels.len(),
+        CHAN * raw_image.rows * raw_image.cols
+    );
+    assert_eq!(cam_model.tstw.len(), CHAN * CHAN);
+    assert_eq!(cam_model.ctrl_pts.len(), cam_model.num_ctrl_pts * CHAN);
+    assert_eq!(cam_model.weights.len(), cam_model.num_ctrl_pts * CHAN);
+    assert_eq!(cam_model.coefs.len(), 4 * CHAN);
+    assert_eq!(cam_model.tonemap.len(), 256 * CHAN);
+
+    let image = HerculesImmBox::from(&raw_image.pixels as &[u8]);
+    let tstw = HerculesImmBox::from(&cam_model.tstw as &[f32]);
+    let ctrl_pts = HerculesImmBox::from(&cam_model.ctrl_pts as &[f32]);
+    let weights = HerculesImmBox::from(&cam_model.weights as &[f32]);
+    let coefs = HerculesImmBox::from(&cam_model.coefs as &[f32]);
+    let tonemap = HerculesImmBox::from(&cam_model.tonemap as &[f32]);
+
+    (
+        raw_image.rows,
+        raw_image.cols,
+        cam_model.num_ctrl_pts,
+        image,
+        tstw,
+        ctrl_pts,
+        weights,
+        coefs,
+        tonemap,
+    )
 }
 
 enum Error {
@@ -154,62 +131,51 @@ pub struct CavaInputs {
 }
 
 pub fn cava_harness(args: CavaInputs) {
-    let CavaInputs {
-        input,
-        output,
-        verify,
-        output_verify,
-        cam_model,
-        crop_rows,
-        crop_cols,
-    } = args;
-    let RawImage { rows, cols, pixels } =
-        read_raw(input, crop_rows, crop_cols).expect("Error loading image");
-    let CamModel {
-        tstw,
-        num_ctrl_pts,
-        ctrl_pts,
-        weights,
-        coefs,
-        tonemap,
-    } = load_cam_model(cam_model, CHAN).expect("Error loading camera model");
+    let (raw_image, cam_model) = make_raw_image_and_cam_model(&args);
+    let (rows, cols, num_ctrl_pts, image, tstw, ctrl_pts, weights, coefs, tonemap) =
+        prepare_hercules_inputs(&raw_image, &cam_model);
+    let mut r = runner!(cava);
 
-    println!(
-        "Running cava with {} rows, {} columns, and {} control points.",
-        rows, cols, num_ctrl_pts
-    );
-    let result = run_cava(
-        rows,
-        cols,
-        num_ctrl_pts,
-        &pixels,
-        &tstw,
-        &ctrl_pts,
-        &weights,
-        &coefs,
-        &tonemap,
-    );
+    let result = async_std::task::block_on(async {
+        HerculesMutBox::from(
+            r.run(
+                rows as u64,
+                cols as u64,
+                num_ctrl_pts as u64,
+                image.to(),
+                tstw.to(),
+                ctrl_pts.to(),
+                weights.to(),
+                coefs.to(),
+                tonemap.to(),
+            )
+            .await,
+        )
+    })
+    .as_slice()
+    .to_vec()
+    .into_boxed_slice();
 
-    if let Some(output) = output {
+    if let Some(output) = args.output {
         extern_image(rows, cols, &*result)
             .save(output)
             .expect("Error saving image");
     }
 
-    if verify {
+    if args.verify {
         let cpu_result = cava_rust::cava(
             rows,
             cols,
             num_ctrl_pts,
-            &pixels,
-            &tstw,
-            &ctrl_pts,
-            &weights,
-            &coefs,
-            &tonemap,
+            &raw_image.pixels,
+            &cam_model.tstw,
+            &cam_model.ctrl_pts,
+            &cam_model.weights,
+            &cam_model.coefs,
+            &cam_model.tonemap,
         );
 
-        if let Some(output) = output_verify {
+        if let Some(output) = args.output_verify {
             extern_image(rows, cols, &cpu_result)
                 .save(output)
                 .expect("Error saving verification image");
@@ -227,5 +193,6 @@ pub fn cava_harness(args: CavaInputs) {
             "Verification failed: maximum pixel difference of {} exceeds threshold of 3",
             max_diff
         );
+        println!("Verified!");
     }
 }