nvtiff-sys 0.5.0

/// Integration tests on the nvTIFF
/// [Decode API](https://docs.nvidia.com/cuda/nvtiff/userguide.html#decode-api-reference)
/// and
/// [Decode Parameters API](https://docs.nvidia.com/cuda/nvtiff/apireference.html#decode-parameters-api)
use std::sync::Arc;

use cudarc::driver::{CudaContext, CudaSlice, CudaStream, DevicePtrMut, DriverError};
use nvtiff_sys::{
    nvtiffDecode, nvtiffDecodeCheckSupported, nvtiffDecodeParams, nvtiffDecodeParamsCreate,
    nvtiffDecodeParamsSetRegions, nvtiffDecodeRegion, nvtiffDecoder, nvtiffDecoderCreateSimple,
    nvtiffImage, nvtiffStatus, nvtiffStream, nvtiffStreamGetNextIFDOffset,
    nvtiffStreamOpenFromFile,
};

#[test]
fn decode_image() -> Result<(), DriverError> {
    // Step 1: Create stream pointers/handles on host (CPU) and device (GPU)

    // Step 1a: Init TIFF stream on host (CPU)
    let mut host_stream = std::mem::MaybeUninit::uninit();
    let mut tiff_stream: *mut nvtiffStream = host_stream.as_mut_ptr();

    // Step 1b: Init CUDA stream on device (GPU)
    let ctx: Arc<CudaContext> = cudarc::driver::CudaContext::new(0)?; // Set on GPU:0
    let stream: Arc<CudaStream> = ctx.default_stream();
    let cuda_stream: *mut nvtiff_sys::CUstream_st = stream.cu_stream().cast::<_>();

    // Step 1c: Init decoder handle
    let mut decoder_handle = std::mem::MaybeUninit::uninit();
    let mut nvtiff_decoder: *mut nvtiffDecoder = decoder_handle.as_mut_ptr();

    let status_decoder: u32 =
        unsafe { nvtiffDecoderCreateSimple(&raw mut nvtiff_decoder, cuda_stream) };
    dbg!(status_decoder);
    assert_eq!(status_decoder, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 2: Parse the TIFF file from disk
    let tiff_cstr = std::ffi::CString::new("images/float32.tif").expect("CString::new failed");
    let tiff_path: *const std::os::raw::c_char = tiff_cstr.as_ptr();

    let status_parse: u32 = unsafe { nvtiffStreamOpenFromFile(tiff_path, &raw mut tiff_stream) };
    dbg!(status_parse);
    assert_eq!(status_parse, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 3a: Create instance of decode parameters handle
    let mut params = std::mem::MaybeUninit::uninit();
    let mut decode_params: *mut nvtiffDecodeParams = params.as_mut_ptr();
    let status_param: u32 = unsafe { nvtiffDecodeParamsCreate(&raw mut decode_params) };
    dbg!(status_param);
    assert_eq!(status_param, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 3b: Allocate memory on device, create nvtiffImage struct instance
    let image_size: usize = 3 * 2 * 4; // Width:3, Height:2, 4 bytes per f32 num
    let mut image_stream: CudaSlice<u8> = stream.alloc_zeros::<u8>(image_size)?;
    let image_ptr: u64 = image_stream.device_ptr_mut(&stream).0;
    let mut plane_data = image_ptr as *mut std::ffi::c_void;
    let mut image_out = nvtiffImage {
        plane_data: &raw mut plane_data,
        plane_pitch_bytes: &mut 0,
        num_planes: 1,
    };

    // Step 3b: Check if image is supported first
    let status_check: u32 = unsafe {
        nvtiffDecodeCheckSupported(
            tiff_stream,
            nvtiff_decoder,
            decode_params,
            &raw const image_out,
        )
    };
    dbg!(status_check);
    assert_eq!(status_check, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 3d: Do the TIFF decoding, defaults to one full image at IFD offset 0
    let status_decode: u32 = unsafe {
        nvtiffDecode(
            tiff_stream,
            nvtiff_decoder,
            decode_params,
            &raw mut image_out,
            cuda_stream,
        )
    };
    dbg!(status_decode); // 2: NVTIFF_STATUS_INVALID_PARAMETER; 8: NVTIFF_STATUS_INTERNAL_ERROR
    assert_eq!(status_decode, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 2c: Transfer decoded bytes from device to host, and check results
    let mut image_out_h: Vec<u8> = vec![0; image_size];
    stream.memcpy_dtoh(&image_stream, &mut image_out_h)?;
    dbg!(image_out_h.clone());
    let float_array: Vec<f32> = image_out_h
        // https://stackoverflow.com/questions/77388769/convert-vecu8-to-vecfloat-in-rust
        // .array_chunks::<4>()
        // .copied()
        .chunks_exact(4)
        .map(TryInto::try_into)
        .map(Result::unwrap)
        .map(f32::from_le_bytes)
        .collect();
    assert_eq!(float_array, vec![1.41, 1.23, 0.78, 0.32, -0.23, -1.88]);

    Ok(())
}

#[test]
fn decode_region() -> Result<(), DriverError> {
    // Step 1: Create stream pointers/handles on host (CPU) and device (GPU)

    // Step 1a: Init TIFF stream on host (CPU)
    let mut host_stream = std::mem::MaybeUninit::uninit();
    let mut tiff_stream: *mut nvtiffStream = host_stream.as_mut_ptr();

    // Step 1b: Init CUDA stream on device (GPU)
    let ctx: Arc<CudaContext> = cudarc::driver::CudaContext::new(0)?; // Set on GPU:0
    let stream: Arc<CudaStream> = ctx.default_stream();
    let cuda_stream: *mut nvtiff_sys::CUstream_st = stream.cu_stream().cast::<_>();

    // Step 1c: Init decoder handle
    let mut decoder_handle = std::mem::MaybeUninit::uninit();
    let mut nvtiff_decoder: *mut nvtiffDecoder = decoder_handle.as_mut_ptr();

    let status_decoder: u32 =
        unsafe { nvtiffDecoderCreateSimple(&raw mut nvtiff_decoder, cuda_stream) };
    dbg!(status_decoder);
    assert_eq!(status_decoder, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 2: Parse the TIFF file from disk
    let tiff_cstr = std::ffi::CString::new("images/cog_dateline.tif").expect("CString::new failed");
    let tiff_path: *const std::os::raw::c_char = tiff_cstr.as_ptr();

    let status_parse: u32 = unsafe { nvtiffStreamOpenFromFile(tiff_path, &raw mut tiff_stream) };
    dbg!(status_parse);
    assert_eq!(status_parse, nvtiffStatus::NVTIFF_STATUS_SUCCESS);
    // let status_print: u32 = unsafe { nvtiffStreamPrint(tiff_stream) };

    // Step 3a: Create instance of decode parameters handle
    let mut params = std::mem::MaybeUninit::uninit();
    let mut decode_params: *mut nvtiffDecodeParams = params.as_mut_ptr();
    let status_param: u32 = unsafe { nvtiffDecodeParamsCreate(&raw mut decode_params) };
    dbg!(status_param);
    assert_eq!(status_param, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 3b: Set decoding region to IFD 1 instead of IFD 0, and an ROI region
    let mut ifd_offset: usize = 0; // start from first IFD
    let status_nextifdoffset: u32 =
        unsafe { nvtiffStreamGetNextIFDOffset(tiff_stream, ifd_offset, &raw mut ifd_offset) };
    dbg!(status_nextifdoffset);
    assert_eq!(status_nextifdoffset, nvtiffStatus::NVTIFF_STATUS_SUCCESS);
    assert_eq!(ifd_offset, 568);

    let region = nvtiffDecodeRegion {
        ifd_offset,
        offset_x: 423,
        offset_y: 390,
        width: 4,
        height: 3,
    };
    let status_region: u32 =
        unsafe { nvtiffDecodeParamsSetRegions(decode_params, &raw const region, 1) };
    dbg!(status_region); // should be 0=SUCCESS
    assert_eq!(status_region, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 3c: Allocate memory on device, create nvtiffImage struct instance
    let image_size: usize = 4 * 3 * 2; // Width:4, Height:3, 2 bytes per u16 num
    let mut image_stream: CudaSlice<u8> = stream.alloc_zeros::<u8>(image_size)?;
    let image_ptr: u64 = image_stream.device_ptr_mut(&stream).0;
    let mut plane_data = image_ptr as *mut std::ffi::c_void;
    let mut image_out = nvtiffImage {
        plane_data: &raw mut plane_data,
        plane_pitch_bytes: &mut 0,
        num_planes: 1,
    };

    // Step 3d: Check if image and ROI bounds is supported first
    let status_check: u32 = unsafe {
        nvtiffDecodeCheckSupported(
            tiff_stream,
            nvtiff_decoder,
            decode_params,
            &raw const image_out,
        )
    };
    dbg!(status_check);
    assert_eq!(status_check, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 3e: Do the TIFF decoding, partial image at IFD 1
    let status_decode: u32 = unsafe {
        nvtiffDecode(
            tiff_stream,
            nvtiff_decoder,
            decode_params,
            &raw mut image_out,
            cuda_stream,
        )
    };
    dbg!(status_decode); // 2: NVTIFF_STATUS_INVALID_PARAMETER; 8: NVTIFF_STATUS_INTERNAL_ERROR
    assert_eq!(status_decode, nvtiffStatus::NVTIFF_STATUS_SUCCESS);

    // Step 2c: Transfer decoded bytes from device to host, and check results
    let mut image_out_h: Vec<u8> = vec![0; image_size];
    stream.memcpy_dtoh(&image_stream, &mut image_out_h)?;
    // dbg!(image_out_h.clone());
    let uint_array: Vec<u16> = image_out_h
        // https://stackoverflow.com/questions/77388769/convert-vecu8-to-vecfloat-in-rust
        // .array_chunks::<2>()
        // .copied()
        .chunks_exact(2)
        .map(TryInto::try_into)
        .map(Result::unwrap)
        .map(u16::from_le_bytes)
        .collect();
    assert_eq!(
        uint_array,
        vec![
            22280, 22280, 22280, 22280, 21890, 22280, 22280, 21952, 21890, 21952, 21952, 21952
        ]
    );

    Ok(())
}