Ускорение 2D БПФ (W) C-смежного, выровненного массива
Я делаю 2D БПФ из 2D массивов комплексных чисел, используя pyFFTW. Эти массивы могут быть очень большими (~128 ГиБ), поэтому время выполнения имеет решающее значение. (Фон - распространение волнового фронта в оптической физике.)
Посмотрите на следующий игрушечный код:
import numpy as np
import pyfftw
import multiprocessing
a = np.random.rand(16384, 16384) + 1j*np.random.rand(16384, 16384)
fft = pyfftw.FFTW(a, a, axes = (0, 1), direction = 'FFTW_FORWARD', flags = ('FFTW_ESTIMATE', 'FFTW_UNALIGNED', 'FFTW_DESTROY_INPUT',), threads = multiprocessing.cpu_count())
a = fft()
FFT занимает несколько секунд для запуска на моей современной 64-битной машине.
Результат и время выполнения остаются одинаковыми при выполнении 2D-БПФ в два этапа (1D-БПФ всех столбцов и всех строк):
fft = pyfftw.FFTW(a, a, axes = (0,), direction = 'FFTW_FORWARD', flags = ('FFTW_ESTIMATE', 'FFTW_UNALIGNED', 'FFTW_DESTROY_INPUT',), threads = multiprocessing.cpu_count())
a = fft()
fft = pyfftw.FFTW(a, a, axes = (1,), direction = 'FFTW_FORWARD', flags = ('FFTW_ESTIMATE', 'FFTW_UNALIGNED', 'FFTW_DESTROY_INPUT',), threads = multiprocessing.cpu_count())
a = fft()
Тем не менее, принимая во внимание время этих шагов по отдельности, вы увидите, что FFT для столбцов примерно в 10 раз медленнее, чем FFT для строк.
Я полагаю, что причина в том, что массив сохраняется в физической памяти строка за строкой. Действительно, флаги дают
C_CONTIGUOUS : True
F_CONTIGUOUS : False
OWNDATA : True
WRITEABLE : True
ALIGNED : True
UPDATEIFCOPY : False
в то время как a.strides дает
(262144, 16)
Таким образом, массив является C-смежным и, кажется, правильно выровнен. Однако удаление флага "FFTW_UNALIGNED" делает FFT для столбцов примерно в 10 раз медленнее (в то время как FFT для строк становится немного быстрее).
Следовательно, мои вопросы:
Может ли быть что-то не так с выравниванием или доступ к столбцам в 10 раз медленнее, чем к строкам для C-смежных массивов физического ограничения?
РЕДАКТИРОВАТЬ: Действительно, коэффициент 10 кажется слишком большим. Давайте сравним простой доступ на чтение / запись к строкам и столбцам:
a[:,0:16384:2]*=1j
а также
a[0:16384:2,:]*=1j
Умножение столбцов с четным индексом (первый вариант) примерно в 2 раза медленнее, чем умножение строк с четным индексом (второй вариант).
РЕДАКТИРОВАТЬ: точный код, введенный в ipython
In [1]: import pyfftw
In [2]: import multiprocessing
In [3]: a = np.random.rand(16384, 16384) + 1j*np.random.rand(16384, 16384)
In [4]: fft = pyfftw.FFTW(a, a, axes = (0, 1), direction = 'FFTW_FORWARD', flags = ('FFTW_ESTIMATE', 'FFTW_UNALIGNED', 'FFTW_DESTROY_INPUT',), threads = multiprocessing.cpu_count())
In [5]: %timeit a = fft()