Skip to content

Commit

Permalink
avfilter/af_afftfilt: add threads support
Browse files Browse the repository at this point in the history
  • Loading branch information
@richardpl
richardpl committed May 24, 2022
1 parent 6076dbc commit d2ef44f
Showing 1 changed file with 92 additions and 47 deletions.
139 changes: 92 additions & 47 deletions libavfilter/af_afftfilt.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -220,59 +229,50 @@ static int config_input(AVFilterLink *inlink)
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;
Expand All @@ -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);
Expand Down Expand Up @@ -360,17 +399,22 @@ 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)
av_freep(&s->fft_out[i]);
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);
Expand Down Expand Up @@ -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,
};

0 comments on commit d2ef44f

Please sign in to comment.