Added optimisation for processing U16x4 images by MulDiv with helps of NEON SIMD instructions.

CHANGELOG.md

@@ -6,6 +6,8 @@
 - Internals of `PixelComponentMapper` changed to use heap to store its data.
 - Fixed link to documentation page in `README.md` file.
 - Added support of optimisation with helps of `NEON SIMD` for convolution of `U16x4` images.
+- Added optimisation for processing `U16x4` images by `MulDiv` with
+  helps of `NEON SIMD` instructions.
 
 ## [2.0.0] - 2022-10-28
 

src/alpha/u16x4/mod.rs

@@ -7,6 +7,8 @@ use super::AlphaMulDiv;
 #[cfg(target_arch = "x86_64")]
 mod avx2;
 mod native;
+#[cfg(target_arch = "aarch64")]
+mod neon;
 #[cfg(target_arch = "x86_64")]
 mod sse4;
 
@@ -21,6 +23,8 @@ impl AlphaMulDiv for U16x4 {
             CpuExtensions::Avx2 => unsafe { avx2::multiply_alpha(src_image, dst_image) },
             #[cfg(target_arch = "x86_64")]
             CpuExtensions::Sse4_1 => unsafe { sse4::multiply_alpha(src_image, dst_image) },
+            // Custom Neon implementation of alpha multiplying are little bit slower than code
+            // generated by Rust compiler for native implementation.
             _ => native::multiply_alpha(src_image, dst_image),
         }
     }
@@ -31,6 +35,8 @@ impl AlphaMulDiv for U16x4 {
             CpuExtensions::Avx2 => unsafe { avx2::multiply_alpha_inplace(image) },
             #[cfg(target_arch = "x86_64")]
             CpuExtensions::Sse4_1 => unsafe { sse4::multiply_alpha_inplace(image) },
+            // Custom Neon implementation of alpha multiplying are little bit slower than code
+            // generated by Rust compiler for native implementation.
             _ => native::multiply_alpha_inplace(image),
         }
     }
@@ -45,6 +51,8 @@ impl AlphaMulDiv for U16x4 {
             CpuExtensions::Avx2 => unsafe { avx2::divide_alpha(src_image, dst_image) },
             #[cfg(target_arch = "x86_64")]
             CpuExtensions::Sse4_1 => unsafe { sse4::divide_alpha(src_image, dst_image) },
+            #[cfg(target_arch = "aarch64")]
+            CpuExtensions::Neon => unsafe { neon::divide_alpha(src_image, dst_image) },
             _ => native::divide_alpha(src_image, dst_image),
         }
     }
@@ -55,6 +63,8 @@ impl AlphaMulDiv for U16x4 {
             CpuExtensions::Avx2 => unsafe { avx2::divide_alpha_inplace(image) },
             #[cfg(target_arch = "x86_64")]
             CpuExtensions::Sse4_1 => unsafe { sse4::divide_alpha_inplace(image) },
+            #[cfg(target_arch = "aarch64")]
+            CpuExtensions::Neon => unsafe { neon::divide_alpha_inplace(image) },
             _ => native::divide_alpha_inplace(image),
         }
     }

src/alpha/u16x4/neon.rs

@@ -0,0 +1,180 @@
+use std::arch::aarch64::*;
+use std::intrinsics::transmute;
+
+use crate::neon_utils;
+use crate::pixels::U16x4;
+use crate::{ImageView, ImageViewMut};
+
+use super::native;
+
+#[target_feature(enable = "neon")]
+pub(crate) unsafe fn multiply_alpha(
+    src_image: &ImageView<U16x4>,
+    dst_image: &mut ImageViewMut<U16x4>,
+) {
+    let src_rows = src_image.iter_rows(0);
+    let dst_rows = dst_image.iter_rows_mut();
+
+    for (src_row, dst_row) in src_rows.zip(dst_rows) {
+        multiply_alpha_row(src_row, dst_row);
+    }
+}
+
+#[target_feature(enable = "neon")]
+pub(crate) unsafe fn multiply_alpha_inplace(image: &mut ImageViewMut<U16x4>) {
+    for dst_row in image.iter_rows_mut() {
+        let src_row = std::slice::from_raw_parts(dst_row.as_ptr(), dst_row.len());
+        multiply_alpha_row(src_row, dst_row);
+    }
+}
+
+#[inline]
+#[target_feature(enable = "neon")]
+unsafe fn multiply_alpha_row(src_row: &[U16x4], dst_row: &mut [U16x4]) {
+    let src_chunks = src_row.chunks_exact(8);
+    let src_remainder = src_chunks.remainder();
+    let mut dst_chunks = dst_row.chunks_exact_mut(8);
+
+    for (src, dst) in src_chunks.zip(&mut dst_chunks) {
+        let mut pixels = neon_utils::load_deintrel_u16x8x4(src, 0);
+        pixels.0 = multiply_color_to_alpha_u16x8(pixels.0, pixels.3);
+        pixels.1 = multiply_color_to_alpha_u16x8(pixels.1, pixels.3);
+        pixels.2 = multiply_color_to_alpha_u16x8(pixels.2, pixels.3);
+        let dst_ptr = dst.as_mut_ptr() as *mut u16;
+        vst4q_u16(dst_ptr, pixels);
+    }
+
+    let src_chunks = src_remainder.chunks_exact(4);
+    let src_remainder = src_chunks.remainder();
+    let dst_reminder = dst_chunks.into_remainder();
+    let mut dst_chunks = dst_reminder.chunks_exact_mut(4);
+    for (src, dst) in src_chunks.zip(&mut dst_chunks) {
+        let mut pixels = neon_utils::load_deintrel_u16x4x4(src, 0);
+        pixels.0 = multiply_color_to_alpha_u16x4(pixels.0, pixels.3);
+        pixels.1 = multiply_color_to_alpha_u16x4(pixels.1, pixels.3);
+        pixels.2 = multiply_color_to_alpha_u16x4(pixels.2, pixels.3);
+        let dst_ptr = dst.as_mut_ptr() as *mut u16;
+        vst4_u16(dst_ptr, pixels);
+    }
+
+    if !src_remainder.is_empty() {
+        let dst_reminder = dst_chunks.into_remainder();
+        native::multiply_alpha_row(src_remainder, dst_reminder);
+    }
+}
+
+#[inline]
+#[target_feature(enable = "neon")]
+unsafe fn multiply_color_to_alpha_u16x8(color: uint16x8_t, alpha: uint16x8_t) -> uint16x8_t {
+    let rounder = vdupq_n_u32(0x8000);
+    let color_lo_u32 = vmlal_u16(rounder, vget_low_u16(color), vget_low_u16(alpha));
+    let color_hi_u32 = vmlal_high_u16(rounder, color, alpha);
+    let color_lo_u32 = vaddhn_u32(color_lo_u32, vshrq_n_u32::<16>(color_lo_u32));
+    let color_hi_u32 = vaddhn_u32(color_hi_u32, vshrq_n_u32::<16>(color_hi_u32));
+    vcombine_u16(color_lo_u32, color_hi_u32)
+}
+
+#[inline]
+#[target_feature(enable = "neon")]
+unsafe fn multiply_color_to_alpha_u16x4(color: uint16x4_t, alpha: uint16x4_t) -> uint16x4_t {
+    let rounder = vdupq_n_u32(0x8000);
+    let color_u32 = vmlal_u16(rounder, color, alpha);
+    vaddhn_u32(color_u32, vshrq_n_u32::<16>(color_u32))
+}
+
+// Divide
+
+#[target_feature(enable = "neon")]
+pub(crate) unsafe fn divide_alpha(
+    src_image: &ImageView<U16x4>,
+    dst_image: &mut ImageViewMut<U16x4>,
+) {
+    let src_rows = src_image.iter_rows(0);
+    let dst_rows = dst_image.iter_rows_mut();
+
+    for (src_row, dst_row) in src_rows.zip(dst_rows) {
+        divide_alpha_row(src_row, dst_row);
+    }
+}
+
+#[target_feature(enable = "neon")]
+pub(crate) unsafe fn divide_alpha_inplace(image: &mut ImageViewMut<U16x4>) {
+    for dst_row in image.iter_rows_mut() {
+        let src_row = std::slice::from_raw_parts(dst_row.as_ptr(), dst_row.len());
+        divide_alpha_row(src_row, dst_row);
+    }
+}
+
+#[target_feature(enable = "neon")]
+pub(crate) unsafe fn divide_alpha_row(src_row: &[U16x4], dst_row: &mut [U16x4]) {
+    let src_chunks = src_row.chunks_exact(8);
+    let src_remainder = src_chunks.remainder();
+    let mut dst_chunks = dst_row.chunks_exact_mut(8);
+
+    for (src, dst) in src_chunks.zip(&mut dst_chunks) {
+        divide_alpha_8_pixels(src, dst);
+    }
+
+    if !src_remainder.is_empty() {
+        let dst_reminder = dst_chunks.into_remainder();
+        let mut src_pixels = [U16x4::new([0; 4]); 8];
+        src_pixels
+            .iter_mut()
+            .zip(src_remainder)
+            .for_each(|(d, s)| *d = *s);
+
+        let mut dst_pixels = [U16x4::new([0; 4]); 8];
+        divide_alpha_8_pixels(src_pixels.as_slice(), dst_pixels.as_mut_slice());
+
+        dst_pixels
+            .iter()
+            .zip(dst_reminder)
+            .for_each(|(s, d)| *d = *s);
+    }
+}
+
+#[inline]
+#[target_feature(enable = "neon")]
+unsafe fn divide_alpha_8_pixels(src: &[U16x4], dst: &mut [U16x4]) {
+    let zero = vdupq_n_u16(0);
+    let alpha_scale = vdupq_n_f32(65535.0);
+    let mut pixels = neon_utils::load_deintrel_u16x8x4(src, 0);
+    let nonzero_alpha_mask = vmvnq_u16(vceqzq_u16(pixels.3));
+
+    // Low
+    let alpha_scaled_u32 = vreinterpretq_u32_u16(vzip1q_u16(pixels.3, zero));
+    let alpha_scaled_f32 = vcvtq_f32_u32(alpha_scaled_u32);
+    let recip_alpha_lo_f32 = vdivq_f32(alpha_scale, alpha_scaled_f32);
+
+    // High
+    let alpha_scaled_u32 = vreinterpretq_u32_u16(vzip2q_u16(pixels.3, zero));
+    let alpha_scaled_f32 = vcvtq_f32_u32(alpha_scaled_u32);
+    let recip_alpha_hi_f32 = vdivq_f32(alpha_scale, alpha_scaled_f32);
+
+    pixels.0 = mul_color_recip_alpha(pixels.0, recip_alpha_lo_f32, recip_alpha_hi_f32, zero);
+    pixels.0 = vandq_u16(pixels.0, nonzero_alpha_mask);
+    pixels.1 = mul_color_recip_alpha(pixels.1, recip_alpha_lo_f32, recip_alpha_hi_f32, zero);
+    pixels.1 = vandq_u16(pixels.1, nonzero_alpha_mask);
+    pixels.2 = mul_color_recip_alpha(pixels.2, recip_alpha_lo_f32, recip_alpha_hi_f32, zero);
+    pixels.2 = vandq_u16(pixels.2, nonzero_alpha_mask);
+
+    let dst_ptr = dst.as_mut_ptr() as *mut u16;
+    vst4q_u16(dst_ptr, pixels);
+}
+
+#[inline]
+#[target_feature(enable = "neon")]
+unsafe fn mul_color_recip_alpha(
+    color: uint16x8_t,
+    recip_alpha_lo: float32x4_t,
+    recip_alpha_hi: float32x4_t,
+    zero: uint16x8_t,
+) -> uint16x8_t {
+    let color_lo_f32 = vcvtq_f32_u32(vreinterpretq_u32_u16(vzip1q_u16(color, zero)));
+    let res_lo_u32 = vcvtaq_u32_f32(vmulq_f32(color_lo_f32, recip_alpha_lo));
+
+    let color_hi_f32 = vcvtq_f32_u32(vreinterpretq_u32_u16(vzip2q_u16(color, zero)));
+    let res_hi_u32 = vcvtaq_u32_f32(vmulq_f32(color_hi_f32, recip_alpha_hi));
+
+    vcombine_u16(vmovn_u32(res_lo_u32), vmovn_u32(res_hi_u32))
+}

src/alpha/u8x4/neon.rs

@@ -137,7 +137,7 @@ unsafe fn divide_alpha_four_pixels(src: &[U8x4], dst: &mut [U8x4]) {
     let zero_alpha_mask = vceqzq_u32(alpha_u32);
     let alpha_f32 = vcvtq_f32_u32(vorrq_u32(alpha_u32, zero_alpha_mask));
     let scaled_alpha_f32 = vdivq_f32(alpha_scale, alpha_f32);
-    let scaled_alpha_u32 = vcvtnq_u32_f32(scaled_alpha_f32);
+    let scaled_alpha_u32 = vcvtaq_u32_f32(scaled_alpha_f32);
     let mma0 = vreinterpretq_u16_u8(vqtbl1q_u8(vreinterpretq_u8_u32(scaled_alpha_u32), shuffle1));
     let mma1 = vreinterpretq_u16_u8(vqtbl1q_u8(vreinterpretq_u8_u32(scaled_alpha_u32), shuffle2));
 

src/neon_utils.rs

@@ -46,6 +46,16 @@ 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_deintrel_u16x4x4<T>(buf: &[T], index: usize) -> uint16x4x4_t {
+    vld4_u16(buf.get_unchecked(index..).as_ptr() as *const u16)
+}
+
+#[inline(always)]
+pub unsafe fn load_deintrel_u16x8x4<T>(buf: &[T], index: usize) -> uint16x8x4_t {
+    vld4q_u16(buf.get_unchecked(index..).as_ptr() as *const u16)
+}
+
 #[inline(always)]
 pub unsafe fn load_i32x2<T>(buf: &[T], index: usize) -> int32x2_t {
     vld1_s32(buf.get_unchecked(index..).as_ptr() as *const i32)
@@ -86,6 +96,11 @@ pub unsafe fn load_i64x2<T>(buf: &[T], index: usize) -> int64x2_t {
     vld1q_s64(buf.get_unchecked(index..).as_ptr() as *const i64)
 }
 
+#[inline(always)]
+pub unsafe fn load_u64x1x4<T>(buf: &[T], index: usize) -> uint64x1x4_t {
+    vld1_u64_x4(buf.get_unchecked(index..).as_ptr() as *const u64)
+}
+
 #[inline(always)]
 pub unsafe fn load_i64x2x2<T>(buf: &[T], index: usize) -> int64x2x2_t {
     vld1q_s64_x2(buf.get_unchecked(index..).as_ptr() as *const i64)
@@ -141,3 +156,15 @@ pub unsafe fn mulhi_u16x8(a: uint16x8_t, b: uint16x8_t) -> uint16x8_t {
     let ab7654 = vmull_high_u16(a, b);
     vuzp2q_u16(vreinterpretq_u16_u32(ab3210), vreinterpretq_u16_u32(ab7654))
 }
+
+/// Multiply the packed unsigned 32-bit integers in a and b, producing
+/// intermediate 64-bit integers, and store the high 32 bits of the intermediate
+/// integers in dst.
+#[inline(always)]
+pub unsafe fn mulhi_u32x4(a: uint32x4_t, b: uint32x4_t) -> uint32x4_t {
+    let a3210 = vget_low_u32(a);
+    let b3210 = vget_low_u32(b);
+    let ab3210 = vmull_u32(a3210, b3210);
+    let ab7654 = vmull_high_u32(a, b);
+    vuzp2q_u32(vreinterpretq_u32_u64(ab3210), vreinterpretq_u32_u64(ab7654))
+}

tests/alpha_tests.rs

@@ -283,6 +283,14 @@ mod multiply_alpha_u16x4 {
         }
     }
 
+    #[cfg(target_arch = "aarch64")]
+    #[test]
+    fn neon_test() {
+        for (s, r) in SRC_PIXELS.into_iter().zip(RES_PIXELS) {
+            mul_div_alpha_test(Oper::Mul, s, r, CpuExtensions::Neon);
+        }
+    }
+
     #[test]
     fn native_test() {
         for (s, r) in SRC_PIXELS.into_iter().zip(RES_PIXELS) {
@@ -489,6 +497,14 @@ mod divide_alpha_u16x4 {
         }
     }
 
+    #[cfg(target_arch = "aarch64")]
+    #[test]
+    fn neon_test() {
+        for (s, r) in SRC_PIXELS.into_iter().zip(RES_PIXELS) {
+            mul_div_alpha_test(OPER, s, r, CpuExtensions::Neon);
+        }
+    }
+
     #[test]
     fn native_test() {
         for (s, r) in SRC_PIXELS.into_iter().zip(RES_PIXELS) {