zune-jpeg 0.5.15

/*
 * Copyright (c) 2025.
 *
 * This software is free software;
 *
 * You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
 */

#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
pub unsafe fn upsample_horizontal_avx2(
    input: &[i16],
    in_near: &[i16],
    in_far: &[i16],
    scratch: &mut [i16],
    output: &mut [i16],
) {
    assert_eq!(input.len() * 2, output.len());
    assert!(input.len() > 2);

    let len = input.len();

    if len < 18 {
        return super::scalar::upsample_horizontal(input, in_near, in_far, scratch, output);
    }

    // First two pixels
    output[0] = input[0];
    output[1] = (input[0] * 3 + input[1] + 2) >> 2;

    let v_three = _mm256_set1_epi16(3);
    let v_two = _mm256_set1_epi16(2);

    let upsample16 = |input: &[i16; 18], output: &mut [i16; 32]| {
        let in_ptr = input.as_ptr();
        let out_ptr = output.as_mut_ptr();

        // SAFETY: The input is 18 * 16 bit long, so the loads are safe.
        let (v_prev, v_curr, v_next) = unsafe {
            (
                _mm256_loadu_si256(in_ptr.add(0) as *const __m256i),
                _mm256_loadu_si256(in_ptr.add(1) as *const __m256i),
                _mm256_loadu_si256(in_ptr.add(2) as *const __m256i),
            )
        };

        let v_common = _mm256_add_epi16(_mm256_mullo_epi16(v_curr, v_three), v_two);

        let v_even = _mm256_srai_epi16(_mm256_add_epi16(v_common, v_prev), 2);
        let v_odd = _mm256_srai_epi16(_mm256_add_epi16(v_common, v_next), 2);

        let v_res_1 = _mm256_unpacklo_epi16(v_even, v_odd);
        let v_res_2 = _mm256_unpackhi_epi16(v_even, v_odd);

        let v_final_1 = _mm256_permute2x128_si256(v_res_1, v_res_2, 0x20);
        let v_final_2 = _mm256_permute2x128_si256(v_res_1, v_res_2, 0x31);

        // SAFETY: The output is 32 * 16 bit long, so the stores are safe.
        unsafe {
            _mm256_storeu_si256(out_ptr as *mut __m256i, v_final_1);
            _mm256_storeu_si256(out_ptr.add(16) as *mut __m256i, v_final_2);
        }
    };

    for (input, output) in input
        .windows(18)
        .step_by(16)
        .zip(output[2..].chunks_exact_mut(32))
    {
        upsample16(input.try_into().unwrap(), output.try_into().unwrap());
    }

    // Upsample the remainder. This may have some overlap, but that's fine.
    if let Some(rest_input) = input.last_chunk::<18>() {
        let end = output.len() - 2;
        if let Some(rest_output) = output[..end].last_chunk_mut::<32>() {
            upsample16(rest_input, rest_output);
        }
    }

    // Last two pixels.
    output[output.len() - 2] = (3 * input[len - 1] + input[len - 2] + 2) >> 2;
    output[output.len() - 1] = input[len - 1];
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
pub unsafe fn upsample_vertical_avx2(
    input: &[i16],
    in_near: &[i16],
    in_far: &[i16],
    scratch: &mut [i16],
    output: &mut [i16],
) {
    assert_eq!(input.len() * 2, output.len());
    assert_eq!(in_near.len(), input.len());
    assert_eq!(in_far.len(), input.len());

    let len = input.len();

    if len < 16 {
        return super::scalar::upsample_vertical(input, in_near, in_far, scratch, output);
    }

    let middle = output.len() / 2;
    let (out_top, out_bottom) = output.split_at_mut(middle);

    let v_three = _mm256_set1_epi16(3);
    let v_two = _mm256_set1_epi16(2);

    let upsample16 = |input: &[i16; 16],
                      in_near: &[i16; 16],
                      in_far: &[i16; 16],
                      out_top: &mut [i16; 16],
                      out_bottom: &mut [i16; 16]| {
        // SAFETY: Inputs are all 16 * 16 bit long, so the loads are safe.
        let (v_in, v_near, v_far) = unsafe {
            (
                _mm256_loadu_si256(input.as_ptr() as *const __m256i),
                _mm256_loadu_si256(in_near.as_ptr() as *const __m256i),
                _mm256_loadu_si256(in_far.as_ptr() as *const __m256i),
            )
        };

        let v_common = _mm256_add_epi16(_mm256_mullo_epi16(v_in, v_three), v_two);

        let v_out_top = _mm256_srai_epi16(_mm256_add_epi16(v_common, v_near), 2);
        let v_out_bottom = _mm256_srai_epi16(_mm256_add_epi16(v_common, v_far), 2);

        // SAFETY: Outputs are 16 * 16 bit long, so the stores are safe.
        unsafe {
            _mm256_storeu_si256(out_top.as_mut_ptr() as *mut __m256i, v_out_top);
            _mm256_storeu_si256(out_bottom.as_mut_ptr() as *mut __m256i, v_out_bottom);
        }
    };

    let chunks = input
        .chunks_exact(16)
        .zip(in_near.chunks_exact(16))
        .zip(in_far.chunks_exact(16))
        .zip(out_top.chunks_exact_mut(16))
        .zip(out_bottom.chunks_exact_mut(16));

    for ((((input, in_near), in_far), out_top), out_bottom) in chunks {
        upsample16(
            input.try_into().unwrap(),
            in_near.try_into().unwrap(),
            in_far.try_into().unwrap(),
            out_top.try_into().unwrap(),
            out_bottom.try_into().unwrap(),
        );
    }

    // Upsample the remainder. This may have some overlap, but that's fine.
    // Edition upgrade will fix this nested awfulness.
    if let Some(rest) = input.last_chunk::<16>() {
        if let Some(rest_near) = in_near.last_chunk::<16>() {
            if let Some(rest_far) = in_far.last_chunk::<16>() {
                if let Some(mut rest_top) = out_top.last_chunk_mut::<16>() {
                    if let Some(mut rest_bottom) = out_bottom.last_chunk_mut::<16>() {
                        upsample16(rest, rest_near, rest_far, &mut rest_top, &mut rest_bottom);
                    }
                }
            }
        }
    }
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
pub unsafe fn upsample_hv_avx2(
    input: &[i16],
    in_near: &[i16],
    in_far: &[i16],
    scratch_space: &mut [i16],
    output: &mut [i16],
) {
    assert_eq!(input.len() * 4, output.len());
    assert!(input.len() * 2 <= scratch_space.len());
    let scratch_space = &mut scratch_space[..input.len() * 2];


    upsample_vertical_avx2(input, in_near, in_far, &mut [], scratch_space);

    let scratch_half = scratch_space.len() / 2;
    let output_half = output.len() / 2;

    let (scratch_top, scratch_bottom) = scratch_space.split_at_mut(scratch_half);
    let (out_top, out_bottom) = output.split_at_mut(output_half);

    let mut t = [0];
    upsample_horizontal_avx2(scratch_top, &[], &[], &mut t, out_top);
    upsample_horizontal_avx2(scratch_bottom, &[], &[], &mut t, out_bottom);
}