Check asm

Author: shssoichiro

src/api/internal.rs

@@ -300,6 +300,7 @@ impl<T: Pixel> ContextInner<T> {
           enc.height,
           enc.bit_depth as u8,
           enc.chroma_sampling,
+          CpuFeatureLevel::default(),
         ))
       } else {
         None

src/denoise.rs renamed to src/denoise/mod.rs

@@ -1,6 +1,8 @@
 use crate::api::FrameQueue;
+use crate::cpu_features::CpuFeatureLevel;
 use crate::EncoderStatus;
 use arrayvec::ArrayVec;
+use cfg_if::cfg_if;
 use ndarray::{Array3, ArrayView3, ArrayViewMut3};
 use ndrustfft::{
   ndfft, ndfft_r2c, ndifft, ndifft_r2c, Complex, FftHandler, R2cFftHandler,
@@ -26,6 +28,9 @@ const COMPLEX_COUNT: usize = (SB_SIZE / 2 + 1) * SB_SIZE * TB_SIZE;
 const CCNT2: usize = COMPLEX_COUNT * 2;
 const INC: usize = SB_SIZE - SO_SIZE;
 
+#[cfg(nasm_x86_64)]
+mod x86;
+
 /// This denoiser is based on the DFTTest plugin from Vapoursynth.
 /// This type of denoising was chosen because it provides
 /// high quality while not being too slow.
@@ -43,6 +48,8 @@ where
   pub(crate) cur_frameno: u64,
 
   denoiser: Box<dyn Denoiser<T>>,
+  #[cfg(feature = "check_asm")]
+  rust_denoiser: DftDenoiserRust<T>,
 }
 
 impl<T> DftDenoiser<T>
@@ -52,7 +59,7 @@ where
   // This should only need to run once per video.
   pub fn new(
     sigma: f32, width: usize, height: usize, bit_depth: u8,
-    chroma_sampling: ChromaSampling,
+    chroma_sampling: ChromaSampling, cpu: CpuFeatureLevel,
   ) -> Self {
     if size_of::<T>() == 1 {
       assert!(bit_depth <= 8);
@@ -81,20 +88,62 @@ where
       effective_heights.push(e_h);
     }
 
-    let denoiser: Box<dyn Denoiser<T>> = Box::new(DftDenoiserRust::new(
+    #[cfg(feature = "check_asm")]
+    let rust_denoiser = DftDenoiserRust::new(
       sigma,
       dest_scale,
       src_scale,
       peak,
-      pad_dimensions,
-      effective_heights,
-    ));
+      pad_dimensions.clone(),
+      effective_heights.clone(),
+    );
+
+    cfg_if! {
+      if #[cfg(nasm_x86_64)] {
+        let denoiser: Box<dyn Denoiser<T>> = match cpu {
+          CpuFeatureLevel::AVX512ICL |
+          CpuFeatureLevel::AVX512 |
+          CpuFeatureLevel::AVX2 => {
+            // SAFETY: This is safe because we verified that AVX2 is enabled.
+            Box::new(unsafe {
+              x86::DftDenoiserAvx2::new(
+                sigma,
+                dest_scale,
+                src_scale,
+                peak,
+                pad_dimensions,
+                effective_heights,
+              )
+            })
+          },
+          _ => Box::new(DftDenoiserRust::new(
+          sigma,
+          dest_scale,
+          src_scale,
+          peak,
+          pad_dimensions,
+          effective_heights,
+        ))
+      };
+      } else {
+        let denoiser: Box<dyn Denoiser<T>> =  Box::new(DftDenoiserRust::new(
+          sigma,
+          dest_scale,
+          src_scale,
+          peak,
+          pad_dimensions,
+          effective_heights,
+        ));
+      }
+    }
 
     DftDenoiser {
       chroma_sampling,
       frame_buffer: VecDeque::with_capacity(TB_MIDPOINT),
       cur_frameno: 0,
       denoiser,
+      #[cfg(feature = "check_asm")]
+      rust_denoiser,
     }
   }
 
@@ -173,6 +222,53 @@ where
     }
     self.denoiser.do_filtering(&pad, &mut dest);
 
+    #[cfg(feature = "check_asm")]
+    {
+      let mut rust_dest = (**orig_frame).clone();
+      let mut pad = ArrayVec::<_, TB_SIZE>::new();
+      for i in 0..TB_SIZE {
+        let dec = self.chroma_sampling.get_decimation().unwrap_or((0, 0));
+        let mut pad_frame = [
+          Plane::new(
+            self.rust_denoiser.pad_dimensions(0).0,
+            self.rust_denoiser.pad_dimensions(0).1,
+            0,
+            0,
+            0,
+            0,
+          ),
+          Plane::new(
+            self.rust_denoiser.pad_dimensions(1).0,
+            self.rust_denoiser.pad_dimensions(1).1,
+            dec.0,
+            dec.1,
+            0,
+            0,
+          ),
+          Plane::new(
+            self.rust_denoiser.pad_dimensions(2).0,
+            self.rust_denoiser.pad_dimensions(2).1,
+            dec.0,
+            dec.1,
+            0,
+            0,
+          ),
+        ];
+
+        let frame = frames.get(&i).unwrap_or(&frames[&TB_MIDPOINT]);
+        self.rust_denoiser.copy_pad(frame, &mut pad_frame);
+        pad.push(pad_frame);
+      }
+      self.rust_denoiser.do_filtering(&pad, &mut rust_dest);
+
+      for p in 0..3 {
+        assert_eq!(
+          dest.planes[p].iter().collect::<Vec<_>>(),
+          rust_dest.planes[p].iter().collect::<Vec<_>>()
+        );
+      }
+    }
+
     if self.frame_buffer.len() == TB_MIDPOINT {
       self.frame_buffer.pop_front();
     }

src/denoise/x86.rs

@@ -0,0 +1,127 @@
+use arrayvec::ArrayVec;
+use ndrustfft::{Complex, FftHandler, R2cFftHandler};
+use v_frame::pixel::Pixel;
+
+use super::{
+  create_window, Denoiser, BLOCK_VOLUME, CCNT2, COMPLEX_COUNT, SB_SIZE,
+  TB_SIZE,
+};
+
+pub(super) struct DftDenoiserAvx2<T>
+where
+  T: Pixel,
+{
+  dest_scale: f32,
+  src_scale: f32,
+  peak: T,
+
+  // These indices refer to planes of the input
+  pad_dimensions: ArrayVec<(usize, usize), 3>,
+  effective_heights: ArrayVec<usize, 3>,
+
+  hw: [f32; BLOCK_VOLUME],
+  dftgc: [Complex<f32>; COMPLEX_COUNT],
+  fft: (R2cFftHandler<f32>, FftHandler<f32>, FftHandler<f32>),
+  sigmas: [f32; CCNT2],
+}
+
+impl<T> DftDenoiserAvx2<T>
+where
+  T: Pixel,
+{
+  #[target_feature(enable = "avx2")]
+  pub unsafe fn new(
+    sigma: f32, dest_scale: f32, src_scale: f32, peak: T,
+    pad_dimensions: ArrayVec<(usize, usize), 3>,
+    effective_heights: ArrayVec<usize, 3>,
+  ) -> Self {
+    let hw = create_window();
+    let mut dftgr = [0f32; BLOCK_VOLUME];
+
+    let fft = (
+      R2cFftHandler::new(SB_SIZE),
+      FftHandler::new(SB_SIZE),
+      FftHandler::new(TB_SIZE),
+    );
+
+    let mut wscale = 0.0f32;
+    for k in 0..BLOCK_VOLUME {
+      dftgr[k] = 255.0 * hw[k];
+      wscale += hw[k].powi(2);
+    }
+    let wscale = 1.0 / wscale;
+
+    let mut sigmas = [0f32; CCNT2];
+    sigmas.fill(sigma / wscale);
+
+    let mut denoiser = DftDenoiserAvx2 {
+      dest_scale,
+      src_scale,
+      peak,
+      pad_dimensions,
+      effective_heights,
+      hw,
+      fft,
+      sigmas,
+      dftgc: [Complex::default(); COMPLEX_COUNT],
+    };
+
+    let mut dftgc = [Complex::default(); COMPLEX_COUNT];
+    denoiser.real_to_complex_3d(&dftgr, &mut dftgc);
+    denoiser.dftgc = dftgc;
+
+    denoiser
+  }
+}
+
+impl<T> Denoiser<T> for DftDenoiserAvx2<T>
+where
+  T: Pixel,
+{
+  #[inline(always)]
+  fn pad_dimensions(&self, plane: usize) -> (usize, usize) {
+    self.pad_dimensions[plane]
+  }
+
+  #[inline(always)]
+  fn effective_height(&self, plane: usize) -> usize {
+    self.effective_heights[plane]
+  }
+
+  #[inline(always)]
+  fn src_scale(&self) -> f32 {
+    self.src_scale
+  }
+
+  #[inline(always)]
+  fn dest_scale(&self) -> f32 {
+    self.dest_scale
+  }
+
+  #[inline(always)]
+  fn peak(&self) -> T {
+    self.peak
+  }
+
+  #[inline(always)]
+  fn hw(&self) -> &[f32; BLOCK_VOLUME] {
+    &self.hw
+  }
+
+  #[inline(always)]
+  fn dftgc(&self) -> &[Complex<f32>; COMPLEX_COUNT] {
+    &self.dftgc
+  }
+
+  #[inline(always)]
+  fn sigmas(&self) -> &[f32; CCNT2] {
+    &self.sigmas
+  }
+
+  #[inline(always)]
+  fn ffts(
+    &mut self,
+  ) -> &mut (R2cFftHandler<f32>, FftHandler<f32>, FftHandler<f32>) {
+    &mut self.fft
+  }
+}