From d2ef44fbb16f3c5675b0c9aa0f281618f7a2a921 Mon Sep 17 00:00:00 2001 From: Paul B Mahol Date: Tue, 24 May 2022 08:52:25 +0200 Subject: [PATCH] avfilter/af_afftfilt: add threads support --- libavfilter/af_afftfilt.c | 139 +++++++++++++++++++++++++------------- 1 file 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, };