Мы можем превзойти производительность функции NumPy на целых 50% с помощью математики

Двумерное преобразование Фурье — один из важнейших алгоритмов компьютерной науки этого столетия. Он нашел широкое применение в нашей повседневной жизни — от фильтров Instagram до обработки MP3-файлов.

Наиболее частой реализацией, используемой рядовым пользователем, иногда даже неосознанно, является адаптация из NumPy. Однако, несмотря на популярность, их алгоритм не является самым эффективным. С помощью нескольких простых манипуляций и статьи 2015 года мы обошли алгоритм NumPy по производительности аж на 30-60%. Основная проблема этой реализации заключается в том, что она изначально основана на слабом с точки зрения производительности алгоритме.

По своей сути алгоритм, реализуемый NumPy, является поочередным применением обычного одномерного БПФ (FFT) к двум измерениям, что очевидно не может быть оптимальным решением.

С другой стороны, в 2015 году двое российских ученых предложили свою версию алгоритма, адаптировав идею одномерного преобразования бабочки для двумерных сигналов. В этой статье мы реализовали их базовую концепцию алгоритма, дополнив ее парочкой своих идей.

Написав наивную реализацию алгоритма из статьи, мы сразу перешли к оптимизации, которая представлена ниже:

void _fft2d( /* Квадратная матрица размера N */ ) {
  // базовый случай {
  if (N == 1) return;
  // } базовый случай 

  int n = N >> 1;
    
  /* псевдокод {
  
  ...Здесь создаются 4 временных матрицы...
    
  // X(x, y, i, j)
  // x, y -- индексация по временным подматрицам
  // i, j -- индексация по строкам и столбцам в каждой подматрице
 
  _fft2d(&X(0, 0), root * root, ...);
  _fft2d(&X(0, 1), root * root, ...);
  _fft2d(&X(1, 0), root * root, ...);
  _fft2d(&X(1, 1), root * root, ...);
  } псевдокод */

  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
      auto x00 = X(0, 0, i, j);
      auto x10 = X(1, 0, i, j) * /* W[i] */;
      auto x01 = X(0, 1, i, j) * /* W[j] */;
      auto x11 = X(1, 1, i, j) * /* W[i] * W[j] */;
      X(0, 0, i, j) = x00 + x10 + x01 + x11;
      X(0, 1, i, j) = x00 + x10 - x01 - x11;
      X(1, 0, i, j) = x00 - x10 + x01 - x11;
      X(1, 1, i, j) = x00 - x10 - x01 + x11;
    }
  }
}

Любой рекурсивный алгоритм может быть усовершенствован путем увеличения размера базового случая. Вот как к этому подошли мы:

void _fft2d( /* Квадратная матрица размера N */ ) {
  // базовый случай {
  if (N == 1) return;
  if (N == 2) {
#define Y(y, x) (V[(y)*rowsize + (x)])
    auto x00 = Y(0, 0);
    auto x10 = Y(1, 0);
    auto x01 = Y(0, 1);
    auto x11 = Y(1, 1);
    Y(0, 0) = x00 + x10 + x01 + x11;
    Y(0, 1) = x00 + x10 - x01 - x11;
    Y(1, 0) = x00 - x10 + x01 - x11;
    Y(1, 1) = x00 - x10 - x01 + x11;
    return;
  }
  // } базовый случай 

  // ...
}

Следующий логический шаг заключался в том, чтобы отказаться от создания четырех ненужных временных подматриц на каждом рекурсивном шаге в пользу всего лишь одной. Для этого мы использовали концепцию из статьи algorithmica, модифицировав ее для двумерной матрицы. Данная возможность также позволяет сократить ненужные выделения памяти и увеличить количество обращений к кэшу.

// Вычисление значений для rev_bits[n]
auto revbits = [](size_t *v, size_t n) {
  int lg_n = log2(n);
  forn(i, n) {
    int revi = 0;
    forn(l, lg_n) revi |= ((i >> l) & 1) << (lg_n - l - 1);
    v[i] = revi;
  }
};

size_t *rev_n = new size_t[N], *rev_m = new size_t[M];
revbits(rev_n, N), revbits(rev_m, M);  

// Преобразующая матрица
forn(i, N) {
  int rev_i = rev_n[i];
  forn(j, M) {
    if ((i < rev_i) || ((i == rev_i) && (j < rev_m[j])))
      std::swap(V[i * M + j], V[rev_i * M + rev_m[j]]);
  }
}

Дальнейшей нашей задачей было предварительное вычисление комплексных корней из единицы:

int mxdim = std::max(N, M);
const int lg_dim = log2(mxdim);
auto W = new fft_type[mxdim];
auto rooti = std::polar(1., (inverse ? 2 : -2) * fft::pi / mxdim);

// Вычисление матрицы поиска корней
auto cur_root = rooti;
W[0] = 1;
forn (i, mxdim - 1) W[i + 1] = W[i] * cur_root;

Как нам может быть достаточно такого массива, спросите вы? Отметим, что в наивной реализации на начальном шаге рекурсии мы передаем массив степеней корня. На следующем уровне рекурсии мы также передаем квадрат этого корня(W[2]). На следующем уровне рекурсии мы проходим через тот же массив степеней, но с шагом 2. Из этого наблюдения можно сделать вывод, что на i-м уровне рекурсии шаг, с которым мы будем проходить через массив W, равен 2ⁱ.

На этом этапе у нас получился следующий код:

void _fft2d(
    fft_type *__restrict__ V,
    size_t N,
    size_t rowsize,
    fft_type *__restrict__ W,
    int step
  ) {
  // базовый случай {
  if (N == 1) return;
  if (N == 2) {
#define Y(y, x) (V[(y)*rowsize + (x)])
    auto x00 = Y(0, 0);
    auto x10 = Y(1, 0);
    auto x01 = Y(0, 1);
    auto x11 = Y(1, 1);
    Y(0, 0) = x00 + x10 + x01 + x11;
    Y(0, 1) = x00 + x10 - x01 - x11;
    Y(1, 0) = x00 - x10 + x01 - x11;
    Y(1, 1) = x00 - x10 - x01 + x11;
    return;
  }
  // } базовый случай

  int n = N >> 1;

#define X(y, x, i, j) (V[((y)*n + (i)) * rowsize + ((x)*n) + j])
#define params n, rowsize, W, (step << 1)
  _fft2d(&X(0, 0, 0, 0), params);
  _fft2d(&X(0, 1, 0, 0), params);
  _fft2d(&X(1, 0, 0, 0), params);
  _fft2d(&X(1, 1, 0, 0), params);

  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
      auto x00 = X(0, 0, i, j);
      auto x10 = X(1, 0, i, j) * W[step * i];
      auto x01 = X(0, 1, i, j) * W[step * j];
      auto x11 = X(1, 1, i, j) * W[step * (i + j)];
      X(0, 0, i, j) = x00 + x10 + x01 + x11;
      X(0, 1, i, j) = x00 + x10 - x01 - x11;
      X(1, 0, i, j) = x00 - x10 + x01 - x11;
      X(1, 1, i, j) = x00 - x10 - x01 + x11;
    }
  }
}

Оригинальный алгоритм имеет еще один очевидный недостаток — он работает только с квадратными матрицами с размерностью, равной степени двойки. С помощью нескольких простых модификаций, мы можем распространить его на прямоугольные матрицы:

void _plan(
    fft_type *__restrict__ V,
    size_t N,
    size_t M,
    size_t rowsize,
    fft_type *__restrict__ W,
    int step_i,
    int step_j
  ) {
  // Вычисление квадратной матрицы
  if (N == M) _fft2d(V, N, rowsize, W, step_i);
  // Выполнение вертикального разбиения
  else if (N > M) {
    int n = N >> 1;
#define Y(y, i, j) (V[((y)*n + (i)) * rowsize + j])
#define params n, M, rowsize, W, (step_i << 1), step_j
    _plan(&Y(0, 0, 0), params);
    _plan(&Y(1, 0, 0), params);

    forn (i, n) {
      forn (j, M) {
        auto y00 = Y(0, i, j);
        auto y10 = Y(1, i, j) * W[i * step_i];
        Y(0, i, j) = y00 + y10;
        Y(1, i, j) = y00 - y10;
      }
    }
  // Выполнение горизонтального разбиения
  } else { /* ...Аналогичный подход... */ }
}

Важно отметить, что NumPy в своем алгоритме делает дополнительное выделение памяти под БПФ, доводя типы в нем до np.complex128; если избежать этого шага, то можно получить выигрыш примерно в 10%. Также мы реализовали многопоточность.

В качестве наглядной демонстрации выигрыша в производительности мы можем представить таблицу с временем выполнения, а также график, показывающий эффективность нашей работы в сравнении с NumPy:

Таблица с результатами

Заключение

Модифицированный алгоритм российских математиков по эффективности обгоняет "строки и столбцы" под капотом NumPy. Несколько логически закономерных манипуляций, таких как увеличение базового случая, значительно повысили нашу производительность.

Очень важно отметить, что те действия, которые мы выполняли при реализации этого алгоритма, могут быть использованы и в случае других алгоритмов, что может пригодиться вам в будущем. Стоит также сказать, что, хотя мы уже приложили значительные усилия, реализация все еще может быть улучшена путем добавления матриц подстановки различных размеров. Цель данной статьи — поделиться исходным кодом, который может помочь улучшить вычисление преобразований в ваших проектах.

Ссылку на репозиторий можно найти ниже, также можно импортировать пакет напрямую из терминала:

pip3 install git+https://github.com/2D-FFT-Project/2d-fft

Спасибо за внимание :)

Репозиторий с исходным кодом

Ресурсы

• FFT, Algorithmica

• Two-dimensional fast Fourier transform: Butterfly in analog of Cooley-Tukey algorithm, by V. S. Tutatchikov for IEEE, 2016

Материал подготовлен в преддверии старта курса OTUS «Алгоритмы и структуры данных». В рамках курсов каждый день проходят открытые уроки по всем IT-направлениям, все темы можно посмотреть в календаре.