use j2k::{
decode_tiles_into, decode_tiles_region_scaled_into, TileBatchOptions, TileDecodeJob,
TileRegionScaledDecodeJob,
};
use j2k_core::{
checked_surface_len, BackendKind, BackendRequest, BatchDecodeError, BatchInfrastructureError,
BufferError, PixelFormat, DEFAULT_MAX_HOST_ALLOCATION_BYTES,
};
use crate::{Error, J2kDecoder, Storage, Surface, SurfaceResidency};
use super::{batch_scheduler_invariant, BatchOp, QueuedRequest};
pub(super) fn decode_cpu_host_batch(
requests: &[QueuedRequest],
) -> Option<Result<Vec<Surface>, Error>> {
decode_cpu_full_batch(requests).or_else(|| decode_cpu_region_scaled_batch(requests))
}
fn decode_cpu_full_batch(requests: &[QueuedRequest]) -> Option<Result<Vec<Surface>, Error>> {
let first = requests.first()?;
if requests.len() <= 1
|| !requests
.iter()
.all(|request| is_cpu_host_full_batch_candidate(request) && request.fmt == first.fmt)
{
return None;
}
Some(decode_cpu_full_batch_inner(requests, first.fmt))
}
fn is_cpu_host_full_batch_candidate(request: &QueuedRequest) -> bool {
matches!(request.op, BatchOp::Full)
&& matches!(request.backend, BackendRequest::Cpu | BackendRequest::Auto)
}
fn decode_cpu_full_batch_inner(
requests: &[QueuedRequest],
fmt: PixelFormat,
) -> Result<Vec<Surface>, Error> {
let mut budget =
crate::batch_allocation::BatchMetadataBudget::new("J2K Metal CPU full batch fallback");
budget.preflight(&[
crate::batch_allocation::BatchMetadataRequest::of::<(u32, u32)>(requests.len()),
crate::batch_allocation::BatchMetadataRequest::of::<(usize, usize)>(requests.len()),
crate::batch_allocation::BatchMetadataRequest::of::<Vec<u8>>(requests.len()),
crate::batch_allocation::BatchMetadataRequest::of::<TileDecodeJob<'_, '_>>(requests.len()),
])?;
let mut dims = budget.try_vec(requests.len(), "J2K Metal CPU full batch dimensions")?;
let mut allocations = budget.try_vec(
requests.len(),
"J2K Metal CPU full batch surface allocations",
)?;
for request in requests {
let decoder = J2kDecoder::new(request.input.as_ref())?;
let tile_dims = decoder.inner.info().dimensions;
let allocation = checked_cpu_batch_surface(tile_dims, fmt)?;
dims.push(tile_dims);
allocations.push(allocation);
}
let mut outputs =
budget.try_vec(allocations.len(), "J2K Metal CPU full batch output owners")?;
for (_, len) in &allocations {
outputs.push(budget.try_filled(*len, 0_u8, "J2K Metal CPU full batch output")?);
}
{
let mut jobs = budget.try_vec(requests.len(), "J2K Metal CPU full batch jobs")?;
for (((request, _dims), (stride, _len)), out) in requests
.iter()
.zip(dims.iter())
.zip(allocations.iter())
.zip(outputs.iter_mut())
{
jobs.push(TileDecodeJob {
input: request.input.as_ref(),
out: out.as_mut_slice(),
stride: *stride,
});
}
decode_tiles_into(&mut jobs, fmt, TileBatchOptions::default()).map_err(cpu_batch_error)?;
}
let mut surfaces = budget.try_vec(requests.len(), "J2K Metal CPU full batch surfaces")?;
for (bytes, dimensions) in outputs.into_iter().zip(dims) {
surfaces.push(host_surface(bytes, dimensions, fmt));
}
Ok(surfaces)
}
fn decode_cpu_region_scaled_batch(
requests: &[QueuedRequest],
) -> Option<Result<Vec<Surface>, Error>> {
let first = requests.first()?;
if requests.len() <= 1
|| !requests.iter().all(|request| {
is_cpu_host_region_scaled_batch_candidate(request) && request.fmt == first.fmt
})
{
return None;
}
Some(decode_cpu_region_scaled_batch_inner(requests, first.fmt))
}
fn is_cpu_host_region_scaled_batch_candidate(request: &QueuedRequest) -> bool {
matches!(request.op, BatchOp::RegionScaled { .. })
&& matches!(request.backend, BackendRequest::Cpu | BackendRequest::Auto)
}
fn decode_cpu_region_scaled_batch_inner(
requests: &[QueuedRequest],
fmt: PixelFormat,
) -> Result<Vec<Surface>, Error> {
let mut budget = crate::batch_allocation::BatchMetadataBudget::new(
"J2K Metal CPU region-scaled batch fallback",
);
budget.preflight(&[
crate::batch_allocation::BatchMetadataRequest::of::<(u32, u32)>(requests.len()),
crate::batch_allocation::BatchMetadataRequest::of::<(usize, usize)>(requests.len()),
crate::batch_allocation::BatchMetadataRequest::of::<Vec<u8>>(requests.len()),
crate::batch_allocation::BatchMetadataRequest::of::<TileRegionScaledDecodeJob<'_, '_>>(
requests.len(),
),
])?;
let mut dims = budget.try_vec(
requests.len(),
"J2K Metal CPU region-scaled batch dimensions",
)?;
let mut allocations = budget.try_vec(
requests.len(),
"J2K Metal CPU region-scaled surface allocations",
)?;
for request in requests {
let BatchOp::RegionScaled { roi, scale } = request.op else {
return Err(batch_scheduler_invariant(
"CPU region-scaled batch contains a non-region-scaled request",
));
};
let dimensions = roi.scaled_covering(scale);
let dims_tuple = (dimensions.w, dimensions.h);
let allocation = checked_cpu_batch_surface(dims_tuple, fmt)?;
dims.push(dims_tuple);
allocations.push(allocation);
}
let mut outputs = budget.try_vec(
allocations.len(),
"J2K Metal CPU region-scaled output owners",
)?;
for (_, len) in &allocations {
outputs.push(budget.try_filled(*len, 0_u8, "J2K Metal CPU region-scaled output")?);
}
{
let mut jobs = budget.try_vec(requests.len(), "J2K Metal CPU region-scaled batch jobs")?;
for ((request, (stride, _len)), out) in requests
.iter()
.zip(allocations.iter())
.zip(outputs.iter_mut())
{
let BatchOp::RegionScaled { roi, scale } = request.op else {
return Err(batch_scheduler_invariant(
"CPU region-scaled job creation received a non-region-scaled request",
));
};
jobs.push(TileRegionScaledDecodeJob {
input: request.input.as_ref(),
out: out.as_mut_slice(),
stride: *stride,
roi,
scale,
});
}
decode_tiles_region_scaled_into(&mut jobs, fmt, TileBatchOptions::default())
.map_err(cpu_batch_error)?;
}
let mut surfaces = budget.try_vec(requests.len(), "J2K Metal CPU region-scaled surfaces")?;
for (bytes, dimensions) in outputs.into_iter().zip(dims) {
surfaces.push(host_surface(bytes, dimensions, fmt));
}
Ok(surfaces)
}
fn checked_cpu_batch_surface(dims: (u32, u32), fmt: PixelFormat) -> Result<(usize, usize), Error> {
checked_surface_len(
dims,
fmt.bytes_per_pixel(),
DEFAULT_MAX_HOST_ALLOCATION_BYTES,
"j2k Metal CPU batch fallback surface",
)
.map_err(Error::from)
}
fn cpu_batch_error(error: j2k::TileBatchError) -> Error {
match error {
BatchDecodeError::Tile(error) => Error::Decode(error.source),
BatchDecodeError::Infrastructure(error) => cpu_batch_infrastructure_error(error),
other => Error::MetalKernel {
message: format!("J2K CPU batch failed: {other}"),
},
}
}
fn cpu_batch_infrastructure_error(error: BatchInfrastructureError) -> Error {
match error {
BatchInfrastructureError::AllocationTooLarge {
what,
requested,
cap,
} => Error::Buffer(BufferError::AllocationTooLarge {
requested,
cap,
what,
}),
BatchInfrastructureError::HostAllocationFailed { what, bytes } => {
Error::Buffer(BufferError::HostAllocationFailed { bytes, what })
}
other => Error::BatchInfrastructure(other),
}
}
fn host_surface(bytes: Vec<u8>, dimensions: (u32, u32), fmt: PixelFormat) -> Surface {
Surface {
backend: BackendKind::Cpu,
residency: SurfaceResidency::Host,
dimensions,
fmt,
pitch_bytes: dimensions.0 as usize * fmt.bytes_per_pixel(),
byte_offset: 0,
storage: Storage::from_host(bytes),
}
}
#[cfg(test)]
mod tests {
use super::*;
use j2k_core::CodecError;
#[test]
fn cpu_batch_infrastructure_preserves_resource_categories() {
let allocation =
cpu_batch_infrastructure_error(BatchInfrastructureError::AllocationTooLarge {
what: "test batch owner",
requested: 9,
cap: 8,
});
assert!(matches!(
&allocation,
Error::Buffer(BufferError::AllocationTooLarge {
what: "test batch owner",
requested: 9,
cap: 8,
})
));
assert!(allocation.is_buffer_error());
let host = cpu_batch_infrastructure_error(BatchInfrastructureError::HostAllocationFailed {
what: "test batch owner",
bytes: 7,
});
assert!(matches!(
&host,
Error::Buffer(BufferError::HostAllocationFailed {
what: "test batch owner",
bytes: 7,
})
));
assert!(host.is_buffer_error());
let source = BatchInfrastructureError::MissingResult { index: 3 };
let scheduler = cpu_batch_infrastructure_error(source);
assert!(matches!(
&scheduler,
Error::BatchInfrastructure(stored) if *stored == source
));
assert!(!scheduler.is_buffer_error());
}
}