Try to implement F32 support for NEON

src/convolution/vertical_f32/mod.rs

@@ -5,8 +5,8 @@ use crate::{CpuExtensions, ImageView, ImageViewMut};
 #[cfg(target_arch = "x86_64")]
 pub(crate) mod avx2;
 pub(crate) mod native;
-// #[cfg(target_arch = "aarch64")]
-// mod neon;
+#[cfg(target_arch = "aarch64")]
+mod neon;
 #[cfg(target_arch = "x86_64")]
 pub(crate) mod sse4;
 // #[cfg(target_arch = "wasm32")]
@@ -28,8 +28,8 @@ pub(crate) fn vert_convolution_f32<T: InnerPixel<Component = f32>>(
         CpuExtensions::Avx2 => avx2::vert_convolution(src_view, dst_view, offset, coeffs),
         #[cfg(target_arch = "x86_64")]
         CpuExtensions::Sse4_1 => sse4::vert_convolution(src_view, dst_view, offset, coeffs),
-        // #[cfg(target_arch = "aarch64")]
-        // CpuExtensions::Neon => neon::vert_convolution(src_view, dst_view, offset, coeffs),
+        #[cfg(target_arch = "aarch64")]
+        CpuExtensions::Neon => neon::vert_convolution(src_view, dst_view, offset, coeffs),
         // #[cfg(target_arch = "wasm32")]
         // CpuExtensions::Simd128 => wasm32::vert_convolution(src_view, dst_view, offset, coeffs),
         _ => native::vert_convolution(src_view, dst_view, offset, coeffs),

src/convolution/vertical_f32/neon.rs

@@ -0,0 +1,126 @@
+use core::arch::aarch64::*;
+
+use super::native;
+use crate::convolution::{Coefficients, CoefficientsChunk};
+use crate::pixels::InnerPixel;
+use crate::{neon_utils, ImageView, ImageViewMut};
+
+pub(crate) fn vert_convolution<T>(
+    src_view: &impl ImageView<Pixel = T>,
+    dst_view: &mut impl ImageViewMut<Pixel = T>,
+    offset: u32,
+    coeffs: &Coefficients,
+) where
+    T: InnerPixel<Component = f32>,
+{
+    let coefficients_chunks = coeffs.get_chunks();
+    let src_x = offset as usize * T::count_of_components();
+
+    let dst_rows = dst_view.iter_rows_mut(0);
+    for (dst_row, coeffs_chunk) in dst_rows.zip(coefficients_chunks) {
+        unsafe {
+            vert_convolution_into_one_row_f32(src_view, dst_row, src_x, coeffs_chunk);
+        }
+    }
+}
+
+#[target_feature(enable = "neon")]
+unsafe fn vert_convolution_into_one_row_f32<T: InnerPixel<Component = f32>>(
+    src_view: &impl ImageView<Pixel = T>,
+    dst_row: &mut [T],
+    mut src_x: usize,
+    coeffs_chunk: CoefficientsChunk,
+) {
+    let mut dst_f32 = T::components_mut(dst_row);
+
+    let mut dst_chunks = dst_f32.chunks_exact_mut(16);
+    for dst_chunk in &mut dst_chunks {
+        multiply_components_of_rows::<_, 8>(src_view, src_x, coeffs_chunk, dst_chunk);
+        src_x += 16;
+    }
+
+    dst_f32 = dst_chunks.into_remainder();
+    dst_chunks = dst_f32.chunks_exact_mut(8);
+    for dst_chunk in &mut dst_chunks {
+        multiply_components_of_rows::<_, 4>(src_view, src_x, coeffs_chunk, dst_chunk);
+        src_x += 8;
+    }
+
+    dst_f32 = dst_chunks.into_remainder();
+    dst_chunks = dst_f32.chunks_exact_mut(4);
+    if let Some(dst_chunk) = dst_chunks.next() {
+        multiply_components_of_rows::<_, 2>(src_view, src_x, coeffs_chunk, dst_chunk);
+        src_x += 4;
+    }
+
+    dst_f32 = dst_chunks.into_remainder();
+    if !dst_f32.is_empty() {
+        let y_start = coeffs_chunk.start;
+        let coeffs = coeffs_chunk.values;
+        native::convolution_by_f32(src_view, dst_f32, src_x, y_start, coeffs);
+    }
+}
+
+#[inline]
+#[target_feature(enable = "neon")]
+pub(crate) unsafe fn multiply_components_of_rows<
+    T: InnerPixel<Component = f32>,
+    const SUMS_COUNT: usize,
+>(
+    src_view: &impl ImageView<Pixel = T>,
+    src_x: usize,
+    coeffs_chunk: CoefficientsChunk,
+    dst_chunk: &mut [f32],
+) {
+    let mut sums = [vdupq_n_f64(0.); SUMS_COUNT];
+    let y_start = coeffs_chunk.start;
+    let mut coeffs = coeffs_chunk.values;
+    let mut y: u32 = 0;
+    let max_rows = coeffs.len() as u32;
+
+    let coeffs_2 = coeffs.chunks_exact(2);
+    coeffs = coeffs_2.remainder();
+    for (src_rows, two_coeffs) in src_view.iter_2_rows(y_start, max_rows).zip(coeffs_2) {
+        let src_rows = src_rows.map(|row| T::components(row).get_unchecked(src_x..));
+        for (&coeff, src_row) in two_coeffs.iter().zip(src_rows) {
+            multiply_components_of_row(&mut sums, coeff, src_row);
+        }
+        y += 2;
+    }
+
+    if let Some(&coeff) = coeffs.first() {
+        if let Some(s_row) = src_view.iter_rows(y_start + y).next() {
+            let src_row = T::components(s_row).get_unchecked(src_x..);
+            multiply_components_of_row(&mut sums, coeff, src_row);
+        }
+    }
+
+    let mut index = 0;
+    for sum in sums {
+        let sum_f32x2 = vcvt_f32_f64(sum);
+        neon_utils::store_f32x2(dst_chunk, index, sum_f32x2);
+        index += 2;
+    }
+}
+
+#[inline]
+#[target_feature(enable = "neon")]
+unsafe fn multiply_components_of_row<const SUMS_COUNT: usize>(
+    sums: &mut [float32x4_t; SUMS_COUNT],
+    coeff: f64,
+    src_row: &[f32],
+) {
+    let coeff_f64x2 = vdupq_n_f64(coeff);
+    let mut i = 0;
+    while i < SUMS_COUNT {
+        let comp03_f32x4 = neon_utils::load_f32x4(src_row, i * 2);
+
+        let comp01_f64x2 = vcvt_f64_f32(vget_low_f32(comp03_f32x4));
+        sums[i] = vmlaq_f64(sums[i], comp01_f64x2, coeff_f64x2);
+        i += 1;
+
+        let comp23_f64x2 = vcvt_f64_f32(vget_high_f32(comp03_f32x4));
+        sums[i] = vmlaq_f64(sums[i], comp23_f64x2, coeff_f64x2);
+        i += 1;
+    }
+}

src/neon_utils.rs

@@ -102,6 +102,36 @@ pub unsafe fn load_u16x8x4<T>(buf: &[T], index: usize) -> uint16x8x4_t {
     vld1q_u16_x4(buf.get_unchecked(index..).as_ptr() as *const u16)
 }
 
+#[inline(always)]
+pub unsafe fn load_f32x2<T>(buf: &[T], index: usize) -> float32x2_t {
+    vld1_f32(buf.get_unchecked(index..).as_ptr() as *const f32)
+}
+
+#[inline(always)]
+pub unsafe fn store_f32x2<T>(buf: &mut [T], index: usize, v: float32x2_t) {
+    vst1_f32(buf.get_unchecked_mut(index..).as_mut_ptr() as *mut f32, v);
+}
+
+#[inline(always)]
+pub unsafe fn load_f32x4<T>(buf: &[T], index: usize) -> float32x4_t {
+    vld1q_f32(buf.get_unchecked(index..).as_ptr() as *const f32)
+}
+
+#[inline(always)]
+pub unsafe fn store_f32x4<T>(buf: &mut [T], index: usize, v: float32x4_t) {
+    vst1q_f32(buf.get_unchecked_mut(index..).as_mut_ptr() as *mut f32, v);
+}
+
+#[inline(always)]
+pub unsafe fn load_f32x4x4<T>(buf: &[T], index: usize) -> float32x4x4_t {
+    vld1q_f32_x4(buf.get_unchecked(index..).as_ptr() as *const f32)
+}
+
+#[inline(always)]
+pub unsafe fn store_f32x4x4<T>(buf: &mut [T], index: usize, v: float32x4x4_t) {
+    vst1q_f32_x4(buf.get_unchecked_mut(index..).as_mut_ptr() as *mut f32, v);
+}
+
 #[inline(always)]
 pub unsafe fn load_deintrel_u16x1x3<T: InnerPixel<Component = u16>>(
     buf: &[T],