avfilter/af_afftfilt: add threads support

1 files changed, 92 insertions, 47 deletions

diff --git a/libavfilter/af_afftfilt.c b/libavfilter/af_afftfilt.c
index a5d3dfe307..b10fa24b09 100644
--- a/libavfilter/af_afftfilt.c
+++ b/libavfilter/af_afftfilt.c
@@ -35,7 +35,7 @@ typedef struct AFFTFiltContext {
     char *img_str;
     int fft_size;
 
-    AVTXContext *fft, *ifft;
+    AVTXContext **fft, **ifft;
     av_tx_fn  tx_fn, itx_fn;
     AVComplexFloat **fft_in;
     AVComplexFloat **fft_out;
@@ -104,19 +104,28 @@ static int config_input(AVFilterLink *inlink)
     AFFTFiltContext *s = ctx->priv;
     char *saveptr = NULL;
     int ret = 0, ch;
-    float overlap, scale;
+    float overlap, scale = 1.f;
     char *args;
     const char *last_expr = "1";
     int buf_size;
 
     s->channels = inlink->ch_layout.nb_channels;
-    ret = av_tx_init(&s->fft, &s->tx_fn, AV_TX_FLOAT_FFT, 0, s->fft_size, &scale, 0);
-    if (ret < 0)
-        return ret;
+    s->fft  = av_calloc(s->channels, sizeof(*s->fft));
+    s->ifft = av_calloc(s->channels, sizeof(*s->ifft));
+    if (!s->fft || !s->ifft)
+        return AVERROR(ENOMEM);
 
-    ret = av_tx_init(&s->ifft, &s->itx_fn, AV_TX_FLOAT_FFT, 1, s->fft_size, &scale, 0);
-    if (ret < 0)
-        return ret;
+    for (int ch = 0; ch < s->channels; ch++) {
+        ret = av_tx_init(&s->fft[ch], &s->tx_fn, AV_TX_FLOAT_FFT, 0, s->fft_size, &scale, 0);
+        if (ret < 0)
+            return ret;
+    }
+
+    for (int ch = 0; ch < s->channels; ch++) {
+        ret = av_tx_init(&s->ifft[ch], &s->itx_fn, AV_TX_FLOAT_FFT, 1, s->fft_size, &scale, 0);
+        if (ret < 0)
+            return ret;
+    }
 
     s->window_size = s->fft_size;
     buf_size = FFALIGN(s->window_size, av_cpu_max_align());
@@ -220,59 +229,50 @@ fail:
     return ret;
 }
 
-static int filter_frame(AVFilterLink *inlink, AVFrame *in)
+static int tx_channel(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
 {
-    AVFilterContext *ctx = inlink->dst;
-    AVFilterLink *outlink = ctx->outputs[0];
     AFFTFiltContext *s = ctx->priv;
-    const int window_size = s->window_size;
-    const float *window_lut = s->window_func_lut;
-    const float f = sqrtf(1.f - s->overlap);
-    double values[VAR_VARS_NB];
-    int ch, n, ret, i;
-    AVFrame *out;
+    const int channels = s->channels;
+    const int start = (channels * jobnr) / nb_jobs;
+    const int end = (channels * (jobnr+1)) / nb_jobs;
 
-    for (ch = 0; ch < inlink->ch_layout.nb_channels; ch++) {
-        const int offset = s->window_size - s->hop_size;
-        float *src = (float *)s->window->extended_data[ch];
+    for (int ch = start; ch < end; ch++) {
         AVComplexFloat *fft_in = s->fft_in[ch];
+        AVComplexFloat *fft_out = s->fft_out[ch];
 
-        memmove(src, &src[s->hop_size], offset * sizeof(float));
-        memcpy(&src[offset], in->extended_data[ch], in->nb_samples * sizeof(float));
-        memset(&src[offset + in->nb_samples], 0, (s->hop_size - in->nb_samples) * sizeof(float));
-
-        for (n = 0; n < window_size; n++) {
-            fft_in[n].re = src[n] * window_lut[n];
-            fft_in[n].im = 0;
-        }
+        s->tx_fn(s->fft[ch], fft_out, fft_in, sizeof(float));
     }
 
-    values[VAR_PTS]         = in->pts;
-    values[VAR_SAMPLE_RATE] = inlink->sample_rate;
-    values[VAR_NBBINS]      = window_size / 2;
-    values[VAR_CHANNELS]    = inlink->ch_layout.nb_channels;
+    return 0;
+}
 
-    for (ch = 0; ch < inlink->ch_layout.nb_channels; ch++) {
-        AVComplexFloat *fft_in = s->fft_in[ch];
-        AVComplexFloat *fft_out = s->fft_out[ch];
+static int filter_channel(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
+{
+    AFFTFiltContext *s = ctx->priv;
+    const int window_size = s->window_size;
+    const float *window_lut = s->window_func_lut;
+    const float f = sqrtf(1.f - s->overlap);
+    const int channels = s->channels;
+    const int start = (channels * jobnr) / nb_jobs;
+    const int end = (channels * (jobnr+1)) / nb_jobs;
+    double values[VAR_VARS_NB];
 
-        s->tx_fn(s->fft, fft_out, fft_in, sizeof(float));
-    }
+    memcpy(values, arg, sizeof(values));
 
-    for (ch = 0; ch < inlink->ch_layout.nb_channels; ch++) {
+    for (int ch = start; ch < end; ch++) {
         AVComplexFloat *fft_out = s->fft_out[ch];
         AVComplexFloat *fft_temp = s->fft_temp[ch];
         float *buf = (float *)s->buffer->extended_data[ch];
-        int x;
+
         values[VAR_CHANNEL] = ch;
 
         if (ctx->is_disabled) {
-            for (n = 0; n < window_size; n++) {
+            for (int n = 0; n < window_size; n++) {
                 fft_temp[n].re = fft_out[n].re;
                 fft_temp[n].im = fft_out[n].im;
             }
         } else {
-            for (n = 0; n <= window_size / 2; n++) {
+            for (int n = 0; n <= window_size / 2; n++) {
                 float fr, fi;
 
                 values[VAR_BIN] = n;
@@ -286,20 +286,59 @@ static int filter_frame(AVFilterLink *inlink, AVFrame *in)
                 fft_temp[n].im = fi;
             }
 
-            for (n = window_size / 2 + 1, x = window_size / 2 - 1; n < window_size; n++, x--) {
+            for (int n = window_size / 2 + 1, x = window_size / 2 - 1; n < window_size; n++, x--) {
                 fft_temp[n].re =  fft_temp[x].re;
                 fft_temp[n].im = -fft_temp[x].im;
             }
         }
 
-        s->itx_fn(s->ifft, fft_out, fft_temp, sizeof(float));
+        s->itx_fn(s->ifft[ch], fft_out, fft_temp, sizeof(float));
 
         memmove(buf, buf + s->hop_size, window_size * sizeof(float));
-        for (i = 0; i < window_size; i++) {
+        for (int i = 0; i < window_size; i++)
             buf[i] += fft_out[i].re * window_lut[i] * f;
+    }
+
+    return 0;
+}
+
+static int filter_frame(AVFilterLink *inlink, AVFrame *in)
+{
+    AVFilterContext *ctx = inlink->dst;
+    AVFilterLink *outlink = ctx->outputs[0];
+    AFFTFiltContext *s = ctx->priv;
+    const int window_size = s->window_size;
+    const float *window_lut = s->window_func_lut;
+    double values[VAR_VARS_NB];
+    int ch, n, ret;
+    AVFrame *out;
+
+    for (ch = 0; ch < inlink->ch_layout.nb_channels; ch++) {
+        const int offset = s->window_size - s->hop_size;
+        float *src = (float *)s->window->extended_data[ch];
+        AVComplexFloat *fft_in = s->fft_in[ch];
+
+        memmove(src, &src[s->hop_size], offset * sizeof(float));
+        memcpy(&src[offset], in->extended_data[ch], in->nb_samples * sizeof(float));
+        memset(&src[offset + in->nb_samples], 0, (s->hop_size - in->nb_samples) * sizeof(float));
+
+        for (n = 0; n < window_size; n++) {
+            fft_in[n].re = src[n] * window_lut[n];
+            fft_in[n].im = 0;
         }
     }
 
+    values[VAR_PTS]         = in->pts;
+    values[VAR_SAMPLE_RATE] = inlink->sample_rate;
+    values[VAR_NBBINS]      = window_size / 2;
+    values[VAR_CHANNELS]    = inlink->ch_layout.nb_channels;
+
+    ff_filter_execute(ctx, tx_channel, NULL, NULL,
+                      FFMIN(s->channels, ff_filter_get_nb_threads(ctx)));
+
+    ff_filter_execute(ctx, filter_channel, values, NULL,
+                      FFMIN(s->channels, ff_filter_get_nb_threads(ctx)));
+
     out = ff_get_audio_buffer(outlink, s->hop_size);
     if (!out) {
         ret = AVERROR(ENOMEM);
@@ -360,10 +399,12 @@ static av_cold void uninit(AVFilterContext *ctx)
     AFFTFiltContext *s = ctx->priv;
     int i;
 
-    av_tx_uninit(&s->fft);
-    av_tx_uninit(&s->ifft);
 
     for (i = 0; i < s->channels; i++) {
+        if (s->ifft)
+            av_tx_uninit(&s->ifft[i]);
+        if (s->fft)
+            av_tx_uninit(&s->fft[i]);
         if (s->fft_in)
             av_freep(&s->fft_in[i]);
         if (s->fft_out)
@@ -371,6 +412,9 @@ static av_cold void uninit(AVFilterContext *ctx)
         if (s->fft_temp)
             av_freep(&s->fft_temp[i]);
     }
+
+    av_freep(&s->fft);
+    av_freep(&s->ifft);
     av_freep(&s->fft_in);
     av_freep(&s->fft_out);
     av_freep(&s->fft_temp);
@@ -412,5 +456,6 @@ const AVFilter ff_af_afftfilt = {
     FILTER_SINGLE_SAMPLEFMT(AV_SAMPLE_FMT_FLTP),
     .activate        = activate,
     .uninit          = uninit,
-    .flags           = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL,
+    .flags           = AVFILTER_FLAG_SUPPORT_TIMELINE_INTERNAL |
+                       AVFILTER_FLAG_SLICE_THREADS,
 };