feat: create SIMD-accelerated version of compute_gru function

xiph/rnnoise#191

Author: JellyBrick

src/rnnoise/include/rnnoise/rnn.h

@@ -27,6 +27,8 @@
 #ifndef RNN_H_
 #define RNN_H_
 
+#include "rnnoise.h"
+
 #include "opus_types.h"
 
 #define WEIGHTS_SCALE (1.f/256)
@@ -58,10 +60,16 @@ typedef struct {
 
 typedef struct RNNState RNNState;
 
+int is_avx2_supported();
+
 void compute_dense(const DenseLayer *layer, float *output, const float *input);
 
 void compute_gru(const GRULayer *gru, float *state, const float *input);
 
+#if defined(__AVX2__)
+void compute_gru_avx2(const GRULayer *gru, float *state, const float *input);
+#endif
+
 void compute_rnn(RNNState *rnn, float *gains, float *vad, const float *input);
 
-#endif /* _MLP_H_ */
+#endif /* RNN_H_ */

src/rnnoise/include/rnnoise/rnn_data.h

@@ -1,32 +1,34 @@
-/*This file is automatically generated from a Keras model*/
-
 #ifndef RNN_DATA_H
 #define RNN_DATA_H
 
 #include "rnn.h"
 
-#define INPUT_DENSE_SIZE 24
-extern const DenseLayer input_dense;
+struct RNNModel {
+  int input_dense_size;
+  const DenseLayer *input_dense;
 
-#define VAD_GRU_SIZE 24
-extern const GRULayer vad_gru;
+  int vad_gru_size;
+  const GRULayer *vad_gru;
 
-#define NOISE_GRU_SIZE 48
-extern const GRULayer noise_gru;
+  int noise_gru_size;
+  const GRULayer *noise_gru;
 
-#define DENOISE_GRU_SIZE 96
-extern const GRULayer denoise_gru;
+  int denoise_gru_size;
+  const GRULayer *denoise_gru;
 
-#define DENOISE_OUTPUT_SIZE 22
-extern const DenseLayer denoise_output;
+  int denoise_output_size;
+  const DenseLayer *denoise_output;
 
-#define VAD_OUTPUT_SIZE 1
-extern const DenseLayer vad_output;
+  int vad_output_size;
+  const DenseLayer *vad_output;
+};
 
 struct RNNState {
-  float vad_gru_state[VAD_GRU_SIZE];
-  float noise_gru_state[NOISE_GRU_SIZE];
-  float denoise_gru_state[DENOISE_GRU_SIZE];
+  const RNNModel *model;
+  float *vad_gru_state;
+  float *noise_gru_state;
+  float *denoise_gru_state;
+  void (*compute_gru_fct)(const GRULayer *gru, float *state, const float *input);
 };
 
 

src/rnnoise/include/rnnoise/rnnoise.h

@@ -1,4 +1,5 @@
-/* Copyright (c) 2017 Mozilla */
+/* Copyright (c) 2018 Gregor Richards
+ * Copyright (c) 2017 Mozilla */
 /*
    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions
@@ -24,6 +25,15 @@
    SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
 
+#ifndef RNNOISE_H
+#define RNNOISE_H 1
+
+#include <stdio.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
 #ifndef RNNOISE_EXPORT
 # if defined(WIN32)
 #  if defined(RNNOISE_BUILD) && defined(DLL_EXPORT)
@@ -39,21 +49,66 @@
 #endif
 
 typedef struct DenoiseState DenoiseState;
+typedef struct RNNModel RNNModel;
 
-#ifdef __cplusplus
-extern "C" {
-#endif
-
+/**
+ * Return the size of DenoiseState
+ */
 RNNOISE_EXPORT int rnnoise_get_size();
 
-RNNOISE_EXPORT int rnnoise_init(DenoiseState *st);
+/**
+ * Return the number of samples processed by rnnoise_process_frame at a time
+ */
+RNNOISE_EXPORT int rnnoise_get_frame_size();
 
-RNNOISE_EXPORT DenoiseState *rnnoise_create();
+/**
+ * Initializes a pre-allocated DenoiseState
+ *
+ * If model is NULL the default model is used.
+ *
+ * See: rnnoise_create() and rnnoise_model_from_file()
+ */
+RNNOISE_EXPORT int rnnoise_init(DenoiseState *st, RNNModel *model);
 
+/**
+ * Allocate and initialize a DenoiseState
+ *
+ * If model is NULL the default model is used.
+ *
+ * The returned pointer MUST be freed with rnnoise_destroy().
+ */
+RNNOISE_EXPORT DenoiseState *rnnoise_create(RNNModel *model);
+
+/**
+ * Free a DenoiseState produced by rnnoise_create.
+ *
+ * The optional custom model must be freed by rnnoise_model_free() after.
+ */
 RNNOISE_EXPORT void rnnoise_destroy(DenoiseState *st);
 
+/**
+ * Denoise a frame of samples
+ *
+ * in and out must be at least rnnoise_get_frame_size() large.
+ */
 RNNOISE_EXPORT float rnnoise_process_frame(DenoiseState *st, float *out, const float *in);
 
+/**
+ * Load a model from a file
+ *
+ * It must be deallocated with rnnoise_model_free()
+ */
+RNNOISE_EXPORT RNNModel *rnnoise_model_from_file(FILE *f);
+
+/**
+ * Free a custom model
+ *
+ * It must be called after all the DenoiseStates referring to it are freed.
+ */
+RNNOISE_EXPORT void rnnoise_model_free(RNNModel *model);
+
 #ifdef __cplusplus
 }
 #endif
+
+#endif

src/rnnoise/src/denoise.c

@@ -1,4 +1,5 @@
-/* Copyright (c) 2017 Mozilla */
+/* Copyright (c) 2018 Gregor Richards
+ * Copyright (c) 2017 Mozilla */
 /*
    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions
@@ -55,13 +56,7 @@
 
 #define SQUARE(x) ((x)*(x))
 
-#define SMOOTH_BANDS 1
-
-#if SMOOTH_BANDS
 #define NB_BANDS 22
-#else
-#define NB_BANDS 21
-#endif
 
 #define CEPS_MEM 8
 #define NB_DELTA_CEPS 6
@@ -73,6 +68,11 @@
 #define TRAINING 0
 #endif
 
+
+/* The built-in model, used if no file is given as input */
+extern const struct RNNModel rnnoise_model_orig;
+
+
 static const opus_int16 eband5ms[] = {
 /*0  200 400 600 800  1k 1.2 1.4 1.6  2k 2.4 2.8 3.2  4k 4.8 5.6 6.8  8k 9.6 12k 15.6 20k*/
   0,  1,  2,  3,  4,  5,  6,  7,  8, 10, 12, 14, 16, 20, 24, 28, 34, 40, 48, 60, 78, 100
@@ -100,7 +100,6 @@ struct DenoiseState {
   RNNState rnn;
 };
 
-#if SMOOTH_BANDS
 void compute_band_energy(float *bandE, const kiss_fft_cpx *X) {
   int i;
   float sum[NB_BANDS] = {0};
@@ -165,32 +164,6 @@ void interp_band_gain(float *g, const float *bandE) {
     }
   }
 }
-#else
-void compute_band_energy(float *bandE, const kiss_fft_cpx *X) {
-  int i;
-  for (i=0;i<NB_BANDS;i++)
-  {
-    int j;
-    opus_val32 sum = 0;
-    for (j=0;j<(eband5ms[i+1]-eband5ms[i])<<FRAME_SIZE_SHIFT;j++) {
-      sum += SQUARE(X[(eband5ms[i]<<FRAME_SIZE_SHIFT) + j].r);
-      sum += SQUARE(X[(eband5ms[i]<<FRAME_SIZE_SHIFT) + j].i);
-    }
-    bandE[i] = sum;
-  }
-}
-
-void interp_band_gain(float *g, const float *bandE) {
-  int i;
-  memset(g, 0, FREQ_SIZE);
-  for (i=0;i<NB_BANDS;i++)
-  {
-    int j;
-    for (j=0;j<(eband5ms[i+1]-eband5ms[i])<<FRAME_SIZE_SHIFT;j++)
-      g[(eband5ms[i]<<FRAME_SIZE_SHIFT) + j] = bandE[i];
-  }
-}
-#endif
 
 
 CommonState common;
@@ -287,19 +260,41 @@ int rnnoise_get_size() {
   return sizeof(DenoiseState);
 }
 
-int rnnoise_init(DenoiseState *st) {
+int rnnoise_get_frame_size() {
+  return FRAME_SIZE;
+}
+
+int rnnoise_init(DenoiseState *st, RNNModel *model) {
   memset(st, 0, sizeof(*st));
+  if (model)
+    st->rnn.model = model;
+  else
+    st->rnn.model = &rnnoise_model_orig;
+  st->rnn.vad_gru_state = calloc(sizeof(float), st->rnn.model->vad_gru_size);
+  st->rnn.noise_gru_state = calloc(sizeof(float), st->rnn.model->noise_gru_size);
+  st->rnn.denoise_gru_state = calloc(sizeof(float), st->rnn.model->denoise_gru_size);
+  st->rnn.compute_gru_fct = &compute_gru;
+
+#if defined(__AVX2__)
+  if(is_avx2_supported() == 1) {
+    st->rnn.compute_gru_fct = &compute_gru_avx2;
+  }
+#endif
+
   return 0;
 }
 
-DenoiseState *rnnoise_create() {
+DenoiseState *rnnoise_create(RNNModel *model) {
   DenoiseState *st;
   st = malloc(rnnoise_get_size());
-  rnnoise_init(st);
+  rnnoise_init(st, model);
   return st;
 }
 
 void rnnoise_destroy(DenoiseState *st) {
+  free(st->rnn.vad_gru_state);
+  free(st->rnn.noise_gru_state);
+  free(st->rnn.denoise_gru_state);
   free(st);
 }
 
@@ -424,13 +419,11 @@ static void frame_synthesis(DenoiseState *st, float *out, const kiss_fft_cpx *y)
 }
 
 static void biquad(float *y, float mem[2], const float *x, const float *b, const float *a, int N) {
-  int i;
-  for (i=0;i<N;i++) {
-    float xi, yi;
-    xi = x[i];
-    yi = x[i] + mem[0];
-    mem[0] = mem[1] + (b[0]*(double)xi - a[0]*(double)yi);
-    mem[1] = (b[1]*(double)xi - a[1]*(double)yi);
+  for (int i=0;i<N;i++) {
+    float xi = x[i];
+    float yi = xi + mem[0];
+    mem[0] = mem[1] + (b[0] * xi - a[0] * yi);
+    mem[1] = (b[1] * xi - a[1] * yi);
     y[i] = yi;
   }
 }
@@ -541,20 +534,21 @@ int main(int argc, char **argv) {
   int vad_cnt=0;
   int gain_change_count=0;
   float speech_gain = 1, noise_gain = 1;
-  FILE *f1, *f2, *fout;
+  FILE *f1, *f2;
+  int maxCount;
   DenoiseState *st;
   DenoiseState *noise_state;
   DenoiseState *noisy;
-  st = rnnoise_create();
-  noise_state = rnnoise_create();
-  noisy = rnnoise_create();
+  st = rnnoise_create(NULL);
+  noise_state = rnnoise_create(NULL);
+  noisy = rnnoise_create(NULL);
   if (argc!=4) {
-    fprintf(stderr, "usage: %s <speech> <noise> <output denoised>\n", argv[0]);
+    fprintf(stderr, "usage: %s <speech> <noise> <count>\n", argv[0]);
     return 1;
   }
   f1 = fopen(argv[1], "r");
   f2 = fopen(argv[2], "r");
-  fout = fopen(argv[3], "w");
+  maxCount = atoi(argv[3]);
   for(i=0;i<150;i++) {
     short tmp[FRAME_SIZE];
     fread(tmp, sizeof(short), FRAME_SIZE, f2);
@@ -566,12 +560,11 @@ int main(int argc, char **argv) {
     float Ln[NB_BANDS];
     float features[NB_FEATURES];
     float g[NB_BANDS];
-    float gf[FREQ_SIZE]={1};
     short tmp[FRAME_SIZE];
     float vad=0;
-    float vad_prob;
     float E=0;
-    if (count==50000000) break;
+    if (count==maxCount) break;
+    if ((count%1000)==0) fprintf(stderr, "%d\r", count);
     if (++gain_change_count > 2821) {
       speech_gain = pow(10., (-40+(rand()%60))/20.);
       noise_gain = pow(10., (-30+(rand()%50))/20.);
@@ -646,37 +639,16 @@ int main(int argc, char **argv) {
       if (vad==0 && noise_gain==0) g[i] = -1;
     }
     count++;
-#if 0
-    for (i=0;i<NB_FEATURES;i++) printf("%f ", features[i]);
-    for (i=0;i<NB_BANDS;i++) printf("%f ", g[i]);
-    for (i=0;i<NB_BANDS;i++) printf("%f ", Ln[i]);
-    printf("%f\n", vad);
-#endif
 #if 1
     fwrite(features, sizeof(float), NB_FEATURES, stdout);
     fwrite(g, sizeof(float), NB_BANDS, stdout);
     fwrite(Ln, sizeof(float), NB_BANDS, stdout);
     fwrite(&vad, sizeof(float), 1, stdout);
-#endif
-#if 0
-    compute_rnn(&noisy->rnn, g, &vad_prob, features);
-    interp_band_gain(gf, g);
-#if 1
-    for (i=0;i<FREQ_SIZE;i++) {
-      X[i].r *= gf[i];
-      X[i].i *= gf[i];
-    }
-#endif
-    frame_synthesis(noisy, xn, X);
-
-    for (i=0;i<FRAME_SIZE;i++) tmp[i] = xn[i];
-    fwrite(tmp, sizeof(short), FRAME_SIZE, fout);
 #endif
   }
   fprintf(stderr, "matrix size: %d x %d\n", count, NB_FEATURES + 2*NB_BANDS + 1);
   fclose(f1);
   fclose(f2);
-  fclose(fout);
   return 0;
 }