From a8da5cf730b70d778404516762047cdc3b24e5fc Mon Sep 17 00:00:00 2001 From: Chin Yun Yu Date: Mon, 7 Jun 2021 11:08:55 +0800 Subject: [PATCH] feat: replace for loop with parallel_for --- torchaudio/csrc/lfilter.cpp | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/torchaudio/csrc/lfilter.cpp b/torchaudio/csrc/lfilter.cpp index dcc5dbe442..84a1fd84e2 100644 --- a/torchaudio/csrc/lfilter.cpp +++ b/torchaudio/csrc/lfilter.cpp @@ -15,18 +15,21 @@ void host_lfilter_core_loop( scalar_t* output_data = padded_output_waveform.data_ptr(); const scalar_t* input_data = input_signal_windows.data_ptr(); const scalar_t* a_coeff_flipped_data = a_coeff_flipped.data_ptr(); - for (int64_t i_channel = 0; i_channel < n_channel; i_channel++) { - for (int64_t i_sample = 0; i_sample < n_samples_input; i_sample++) { - int64_t offset_input = i_channel * n_samples_input; - int64_t offset_output = i_channel * n_samples_output; - scalar_t a0 = input_data[offset_input + i_sample]; - for (int64_t i_coeff = 0; i_coeff < n_order; i_coeff++) { - a0 -= output_data[offset_output + i_sample + i_coeff] * - a_coeff_flipped_data[i_coeff]; + + at::parallel_for(0, n_channel, 1, [&](int64_t begin, int64_t end) { + for (auto i_channel = begin; i_channel < end; i_channel++) { + for (int64_t i_sample = 0; i_sample < n_samples_input; i_sample++) { + int64_t offset_input = i_channel * n_samples_input; + int64_t offset_output = i_channel * n_samples_output; + scalar_t a0 = input_data[offset_input + i_sample]; + for (int64_t i_coeff = 0; i_coeff < n_order; i_coeff++) { + a0 -= output_data[offset_output + i_sample + i_coeff] * + a_coeff_flipped_data[i_coeff]; + } + output_data[offset_output + i_sample + n_order - 1] = a0; } - output_data[offset_output + i_sample + n_order - 1] = a0; } - } + }); } void cpu_lfilter_core_loop(