Improve codec batching and Metal kernels

Cargo.toml

@@ -111,7 +111,7 @@ jpeg-encoder = "0.7.0"
 jpeg-decoder = { version = "0.3.2", default-features = false }
 zune-core = "0.5.1"
 zune-jpeg = "0.5.15"
-fearless_simd = "0.7"
+fearless_simd = "1.0"
 libm = { version = "0.2", default-features = false }
 openjpeg-sys = "1.0.12"
 proc-macro2 = "1.0.106"

crates/j2k-compare/src/common.rs

@@ -432,7 +432,42 @@ pub fn mib_per_second(bytes: usize, elapsed_us: f64) -> f64 {
 mod tests {
     use std::path::Path;
 
-    use super::infer_corpus_category;
+    use super::{batch_size_config_from_values, infer_corpus_category, BatchSizeConfig};
+
+    #[test]
+    fn batch_size_config_preserves_defaults_and_independent_legacy_overrides() {
+        for (case_sizes, mixed_sizes, legacy, expected_case, expected_mixed) in [
+            (None, None, None, vec![1], vec![1, 16, 256, 1024]),
+            (Some("3"), None, Some(vec![2, 4]), vec![3], vec![2, 4]),
+            (
+                None,
+                Some("8,16"),
+                Some(vec![2, 4]),
+                vec![2, 4],
+                vec![8, 16],
+            ),
+        ] {
+            let config = batch_size_config_from_values(
+                case_sizes,
+                mixed_sizes,
+                legacy,
+                "case batch sizes",
+                "mixed batch sizes",
+                &[1],
+                &[1, 16, 256, 1024],
+            )
+            .expect("batch size config parses");
+
+            assert_eq!(
+                config,
+                BatchSizeConfig {
+                    case_batch_sizes: expected_case,
+                    mixed_batch_sizes: expected_mixed,
+                },
+                "case sizes {case_sizes:?}, mixed sizes {mixed_sizes:?}"
+            );
+        }
+    }
 
     #[test]
     fn corpus_category_rules_cover_every_supported_needle() {

crates/j2k-compare/src/encode_compare.rs

@@ -137,52 +137,10 @@ mod tests {
     };
     use super::{
         canonicalize_manifest_row_path, measurement_row, EncoderKind, EncoderTool, ImageCase,
-        Measurement, MetadataInput, MixedImageBatch, DEFAULT_CASE_BATCH_SIZES,
-        DEFAULT_MIXED_BATCH_SIZES,
+        Measurement, MetadataInput, MixedImageBatch,
     };
-    use crate::common;
     use std::path::Path;
 
-    fn test_batch_size_config_from_values(
-        case_batch_sizes: Option<&str>,
-        mixed_batch_sizes: Option<&str>,
-        legacy: Option<Vec<usize>>,
-    ) -> Result<common::BatchSizeConfig, String> {
-        common::batch_size_config_from_values(
-            case_batch_sizes,
-            mixed_batch_sizes,
-            legacy,
-            "J2K_ENCODE_COMPARE_CASE_BATCH_SIZES",
-            "J2K_ENCODE_COMPARE_MIXED_BATCH_SIZES",
-            DEFAULT_CASE_BATCH_SIZES,
-            DEFAULT_MIXED_BATCH_SIZES,
-        )
-    }
-
-    #[test]
-    fn encode_batch_config_defaults_keep_large_batches_mixed_only() {
-        let config = test_batch_size_config_from_values(None, None, None)
-            .expect("default batch config parses");
-
-        assert_eq!(config.case_batch_sizes, DEFAULT_CASE_BATCH_SIZES);
-        assert_eq!(config.mixed_batch_sizes, DEFAULT_MIXED_BATCH_SIZES);
-    }
-
-    #[test]
-    fn encode_batch_config_split_env_overrides_legacy_independently() {
-        let config = test_batch_size_config_from_values(Some("3"), None, Some(vec![2, 4]))
-            .expect("case override with legacy config parses");
-
-        assert_eq!(config.case_batch_sizes, vec![3]);
-        assert_eq!(config.mixed_batch_sizes, vec![2, 4]);
-
-        let config = test_batch_size_config_from_values(None, Some("8,16"), Some(vec![2, 4]))
-            .expect("mixed override with legacy config parses");
-
-        assert_eq!(config.case_batch_sizes, vec![2, 4]);
-        assert_eq!(config.mixed_batch_sizes, vec![8, 16]);
-    }
-
     #[test]
     fn encode_manifest_path_remaps_to_supplied_fixture_root_by_suffix() {
         let root = std::env::current_dir()

crates/j2k-compare/src/fixture_compare/tests.rs

@@ -8,52 +8,10 @@ use super::metadata::{
 use super::{
     canonicalize_manifest_row_path, publication_blockers, unique_input_count, BenchmarkMode, Codec,
     Container, DecoderKind, FixtureCase, MixedFixtureBatch, Operation, OperationClass,
-    DEFAULT_CASE_BATCH_SIZES, DEFAULT_MIXED_BATCH_SIZES,
 };
-use crate::common;
 use j2k_core::{Downscale, PixelFormat, Rect};
 use std::path::Path;
 
-fn test_batch_size_config_from_values(
-    case_batch_sizes: Option<&str>,
-    mixed_batch_sizes: Option<&str>,
-    legacy: Option<Vec<usize>>,
-) -> Result<common::BatchSizeConfig, String> {
-    common::batch_size_config_from_values(
-        case_batch_sizes,
-        mixed_batch_sizes,
-        legacy,
-        "J2K_FIXTURE_COMPARE_CASE_BATCH_SIZES",
-        "J2K_FIXTURE_COMPARE_MIXED_BATCH_SIZES",
-        DEFAULT_CASE_BATCH_SIZES,
-        DEFAULT_MIXED_BATCH_SIZES,
-    )
-}
-
-#[test]
-fn decode_batch_config_defaults_keep_large_batches_mixed_only() {
-    let config =
-        test_batch_size_config_from_values(None, None, None).expect("default batch config parses");
-
-    assert_eq!(config.case_batch_sizes, DEFAULT_CASE_BATCH_SIZES);
-    assert_eq!(config.mixed_batch_sizes, DEFAULT_MIXED_BATCH_SIZES);
-}
-
-#[test]
-fn decode_batch_config_split_env_overrides_legacy_independently() {
-    let config = test_batch_size_config_from_values(Some("3"), None, Some(vec![2, 4]))
-        .expect("case override with legacy config parses");
-
-    assert_eq!(config.case_batch_sizes, vec![3]);
-    assert_eq!(config.mixed_batch_sizes, vec![2, 4]);
-
-    let config = test_batch_size_config_from_values(None, Some("8,16"), Some(vec![2, 4]))
-        .expect("mixed override with legacy config parses");
-
-    assert_eq!(config.case_batch_sizes, vec![2, 4]);
-    assert_eq!(config.mixed_batch_sizes, vec![8, 16]);
-}
-
 #[test]
 fn decode_manifest_path_remaps_to_supplied_fixture_root_by_suffix() {
     let root = std::env::current_dir()

crates/j2k-jpeg-metal/Cargo.toml

@@ -26,6 +26,7 @@ j2k-profile = { path = "../j2k-profile", version = "=0.11.1" }
 thiserror = { workspace = true }
 
 [target.'cfg(target_os = "macos")'.dependencies]
+dispatch2 = { version = "=0.3.1", default-features = false, features = ["std"] }
 objc2 = { workspace = true }
 objc2-metal = { workspace = true }
 

crates/j2k-jpeg-metal/README.md

@@ -23,9 +23,12 @@ fast 4:2:0, 4:2:2, or 4:4:4 baseline packet and only `Gray8`, `Rgb8`, or
 unsupported output formats return `UnsupportedMetalRequest` instead of silently
 falling back.
 
-`BackendRequest::Auto` stays conservative. Single-image decode remains CPU even
-when fast-packet capabilities match. Batched and resident-output paths are the
-places to look for Metal wins, and any future Auto widening should be backed by
+`BackendRequest::Auto` keeps single-image decode on the CPU. Full RGB batches
+can use Metal for at least 16 compatible non-restart 4:2:0 or 4:2:2 tiles of at
+least 256×256 pixels, including distinct inputs. The tiles must share dimensions,
+tables, and checkpoint count. Smaller, mixed-table, restart-coded, and scaled
+batches remain on the CPU. This threshold follows completed
+submission measurements on an M4 Pro; further widening should be backed by
 the benchmark groups documented in
 [`docs/routing-benchmarks.md`](docs/routing-benchmarks.md).
 

crates/j2k-jpeg-metal/benches/compare.rs

@@ -608,6 +608,36 @@ fn auto_decode_tile_batch(bytes: &[u8], batch_size: usize) {
     device_decode_tile_batch(bytes, batch_size, BackendRequest::Auto);
 }
 
+/// Decode context, scratch, and Metal session kept across iterations, as a
+/// long-running tile server would, unlike `device_decode_tile_batch`.
+#[derive(Default)]
+struct RetainedTileBatchState {
+    ctx: JpegDecoderContext,
+    pool: ScratchPool,
+    session: MetalSession,
+}
+
+impl RetainedTileBatchState {
+    fn decode(&mut self, bytes: &[u8], batch_size: usize, backend: BackendRequest) {
+        let submissions = (0..batch_size)
+            .map(|_| {
+                <Codec as TileBatchDecodeSubmit>::submit_tile_to_device(
+                    &mut self.ctx,
+                    &mut self.session,
+                    &mut self.pool,
+                    bytes,
+                    PixelFormat::Rgb8,
+                    backend,
+                )
+                .expect("submit")
+            })
+            .collect::<Vec<_>>();
+        for submission in submissions {
+            std::hint::black_box(submission.wait().expect("surface"));
+        }
+    }
+}
+
 fn device_decode_tile_batch(bytes: &[u8], batch_size: usize, backend: BackendRequest) {
     let mut ctx = JpegDecoderContext::default();
     let mut pool = ScratchPool::new();
@@ -630,6 +660,65 @@ fn device_decode_tile_batch(bytes: &[u8], batch_size: usize, backend: BackendReq
     }
 }
 
+#[cfg(target_os = "macos")]
+fn bench_distinct_batch_routing(c: &mut Criterion) {
+    let mut group = c.benchmark_group("jpeg_distinct_batch_routing");
+    for (family, sampling) in [
+        ("420", SamplingFactor::F_2_2),
+        ("422", SamplingFactor::F_2_1),
+    ] {
+        for side in [64_u16, 256] {
+            let inputs = (0..64)
+                .map(|index| generated_rgb_jpeg_variant(side, side, sampling, None, index))
+                .collect::<Vec<_>>();
+            for count in [16, 64] {
+                let decode = |backend| {
+                    let mut context = JpegDecoderContext::default();
+                    let mut session = MetalSession::default();
+                    let mut pool = ScratchPool::new();
+                    let submissions = inputs[..count]
+                        .iter()
+                        .map(|bytes| {
+                            <Codec as TileBatchDecodeSubmit>::submit_tile_to_device(
+                                &mut context,
+                                &mut session,
+                                &mut pool,
+                                bytes,
+                                PixelFormat::Rgb8,
+                                backend,
+                            )
+                            .expect("distinct routing submission")
+                        })
+                        .collect::<Vec<_>>();
+                    submissions
+                        .into_iter()
+                        .map(|submission| submission.wait().expect("distinct routing result"))
+                        .collect::<Vec<_>>()
+                };
+                let metal = decode(BackendRequest::Metal);
+                let cpu = decode(BackendRequest::Cpu);
+                for (actual, expected) in metal.iter().zip(&cpu) {
+                    assert_metal_surface_pixels(
+                        actual,
+                        expected.as_bytes().expect("CPU pixels").as_ref(),
+                    );
+                }
+                for (label, backend) in [
+                    ("cpu", BackendRequest::Cpu),
+                    ("metal", BackendRequest::Metal),
+                    ("auto", BackendRequest::Auto),
+                ] {
+                    group.bench_function(
+                        format!("{family}/{side}x{side}/batch{count}/{label}"),
+                        |b| b.iter(|| std::hint::black_box(decode(backend))),
+                    );
+                }
+            }
+        }
+    }
+    group.finish();
+}
+
 fn metal_decode_tile_batch_scaled(bytes: &[u8], batch_size: usize, factor: Downscale) {
     device_decode_tile_batch_scaled(bytes, batch_size, factor, BackendRequest::Metal);
 }
@@ -885,6 +974,13 @@ fn bench_full_and_tile_decode_groups(c: &mut Criterion, inputs: &[BenchInput], h
             wsi_tile_batch_rgb.bench_function(format!("{}/metal", input.name), |b| {
                 b.iter(|| metal_decode_tile_batch(&input.bytes, 64));
             });
+            let mut retained = RetainedTileBatchState::default();
+            wsi_tile_batch_rgb.bench_function(
+                format!("{}/metal_retained_session", input.name),
+                |b| {
+                    b.iter(|| retained.decode(&input.bytes, 64, BackendRequest::Metal));
+                },
+            );
         }
         wsi_tile_batch_rgb.bench_function(format!("{}/auto", input.name), |b| {
             b.iter(|| auto_decode_tile_batch(&input.bytes, 64));
@@ -1634,6 +1730,8 @@ fn bench_compare(c: &mut Criterion) {
     #[cfg(target_os = "macos")]
     if has_metal {
         distinct_batch::bench(c);
+        distinct_batch::bench_mixed(c);
+        bench_distinct_batch_routing(c);
         representative_matrix::bench(c);
     }
     bench_resident_texture_batches(c, &inputs, has_metal);