j2k-cuda 0.11.0

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Prepared-input validation and CUDA adapter mapping.

#[cfg(feature = "cuda-runtime")]
use super::{
    cuda_range_storage, Arc, BackendKind, CudaResidentBatchBuffer, CudaSurfaceStats,
    PreparedBatchGroup, Surface, SurfaceResidency,
};
use super::{
    BatchColor, BatchGroupInfo, BatchLayout, Error, J2kDecodeWarning, PixelFormat, PreparedImage,
};

pub(super) fn group_pixel_format(info: &BatchGroupInfo) -> Result<PixelFormat, Error> {
    info.native_pixel_format().ok_or(Error::capability_rejected(
        j2k_core::CapabilityRejection::unsupported_format(
            "CUDA batch output color/sample type is unsupported",
        ),
    ))
}

pub(super) fn validate_layout(info: &BatchGroupInfo) -> Result<(), Error> {
    if !matches!(info.layout, BatchLayout::Nchw | BatchLayout::Nhwc) {
        return Err(Error::capability_rejected(
            j2k_core::CapabilityRejection::unsupported_format(
                "CUDA batch output layout is unsupported",
            ),
        ));
    }
    Ok(())
}

#[cfg(feature = "cuda-runtime")]
pub(super) fn native_decode_settings(settings: j2k::DecodeSettings) -> j2k_native::DecodeSettings {
    j2k_native::DecodeSettings {
        resolve_palette_indices: true,
        strict: settings.is_strict(),
        target_resolution: None,
    }
}

#[cfg(feature = "cuda-runtime")]
pub(super) fn native_referenced_htj2k_plan(
    plan: &j2k::PreparedHtj2kPlan,
) -> &j2k_native::J2kReferencedHtj2kPlan {
    plan.geometry()
}

#[cfg(feature = "cuda-runtime")]
pub(super) fn native_referenced_classic_plan(
    plan: &j2k::PreparedClassicPlan,
) -> &j2k_native::J2kReferencedClassicPlan {
    plan.geometry()
}

#[cfg(feature = "cuda-runtime")]
pub(super) fn native_color_inputs(
    group: &PreparedBatchGroup,
) -> Result<Vec<crate::decoder::NativeColorBatchInput<'_>>, Error> {
    group
        .images()
        .iter()
        .zip(group.source_indices().iter().copied())
        .map(|(image, source_index)| {
            let referenced_plan = match image.htj2k_plan() {
                Some(prepared_plan)
                    if group.info().color == BatchColor::Rgb && prepared_plan.is_color() =>
                {
                    Some(native_referenced_htj2k_plan(prepared_plan))
                }
                Some(prepared_plan)
                    if group.info().color == BatchColor::Rgba && prepared_plan.is_rgba() =>
                {
                    Some(native_referenced_htj2k_plan(prepared_plan))
                }
                Some(_) => {
                    return Err(Error::capability_rejected(
                        j2k_core::CapabilityRejection::geometry_mismatch(
                            "exact CUDA color batch received incompatible prepared geometry",
                        ),
                    ))
                }
                None => None,
            };
            let referenced_classic_plan = match image.classic_plan() {
                Some(prepared_plan)
                    if group.info().color == BatchColor::Rgb && prepared_plan.is_color() =>
                {
                    Some(native_referenced_classic_plan(prepared_plan))
                }
                Some(prepared_plan)
                    if group.info().color == BatchColor::Rgba && prepared_plan.is_rgba() =>
                {
                    Some(native_referenced_classic_plan(prepared_plan))
                }
                Some(_) => {
                    return Err(Error::capability_rejected(
                        j2k_core::CapabilityRejection::geometry_mismatch(
                            "exact CUDA color batch received incompatible classic geometry",
                        ),
                    ));
                }
                None => None,
            };
            if referenced_plan.is_none() && referenced_classic_plan.is_none() {
                return Err(Error::capability_rejected(
                    j2k_core::CapabilityRejection::missing_prepared_plan(
                        "exact CUDA color batch requires a supported prepared device plan",
                    ),
                ));
            }
            Ok(crate::decoder::NativeColorBatchInput {
                source_index,
                bytes: image.bytes().as_ref(),
                device_plan: image.plan(),
                referenced_plan,
                referenced_classic_plan,
                settings: native_decode_settings(group.options().settings),
            })
        })
        .collect()
}

#[cfg(feature = "cuda-runtime")]
pub(super) fn native_color_group_storage(
    info: &BatchGroupInfo,
    fmt: PixelFormat,
    output: crate::decoder::NativeColorOwnedBatch,
) -> (Vec<Surface>, CudaResidentBatchBuffer) {
    let crate::decoder::NativeColorOwnedBatch {
        buffer,
        ranges,
        execution,
    } = output;
    let shared = Arc::new(buffer);
    let surfaces = if info.layout == BatchLayout::Nhwc {
        ranges
            .iter()
            .map(|range| Surface {
                backend: BackendKind::Cuda,
                residency: SurfaceResidency::CudaResidentDecode,
                dimensions: info.dimensions,
                fmt,
                pitch_bytes: info.dimensions.0 as usize * fmt.bytes_per_pixel(),
                stats: CudaSurfaceStats {
                    total: execution.kernel_dispatches(),
                    copy: execution.copy_kernel_dispatches(),
                    decode: execution.decode_kernel_dispatches(),
                },
                storage: cuda_range_storage(shared.clone(), range.offset, range.len),
            })
            .collect()
    } else {
        Vec::new()
    };
    (
        surfaces,
        CudaResidentBatchBuffer {
            buffer: shared,
            ranges,
        },
    )
}

pub(super) fn decode_warnings(
    images: &[PreparedImage],
) -> Result<Vec<Vec<J2kDecodeWarning>>, Error> {
    let mut budget = crate::allocation::HostPhaseBudget::new("CUDA batch decode warnings");
    let mut warnings = budget.try_vec_with_capacity(images.len())?;
    for image in images {
        let mut image_warnings = Vec::new();
        if image.used_lenient_metadata_recovery() {
            // `J2kDecodeWarning` is a ZST, so this push does not allocate.
            image_warnings.push(J2kDecodeWarning::LenientMetadataRecovery);
        }
        warnings.push(image_warnings);
    }
    Ok(warnings)
}