Как распараллелить простой цикл над матрицей?
У меня большая матрица, и я хочу вывести все индексы, в которых элементы матрицы меньше 0. У меня есть этот MWE в цифровом формате:
import numba as nb
import numpy as np
A = np.random.random(size = (1000, 1000)) - 0.1
@nb.njit(cache=True)
def numba_only(arr):
rows = np.empty(arr.shape[0]*arr.shape[1])
cols = np.empty(arr.shape[0]*arr.shape[1])
idx = 0
for i in range(arr.shape[0]):
for j in range(A.shape[1]):
if arr[i, j] < 0:
rows[idx] = i
cols[idx] = j
idx += 1
return rows[:idx], cols[:idx]
Время, я получаю:
%timeit numba_only(A)
2.29 ms ± 114 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Это быстрее, чем np.where(A<0) (без цифр), что дает:
%timeit numpy_only(A)
3.56 ms ± 59.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Но это немного медленнее, чем np.where, обернутое @nb.njit:
@nb.njit(cache=True)
def numba_where(arr):
return np.where(arr<0)
который дает:
%timeit numba_where(A)
1.76 ms ± 84.9 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
Можно ли как-то ускорить числовой код за счет распараллеливания? Я понимаю, что это может быть связано с памятью, но я думаю, что современное оборудование должно обеспечивать некоторый уровень параллельного доступа к памяти.
Я не уверен, как использовать nb.prange для достижения этой цели из-за индекса idx в цикле.
Ответы 1
Предоставленный алгоритм по своей сути является последовательным из-за длинной цепочки зависимостей команд, которая также зависит от данных. Однако эквивалентную операцию можно выполнить параллельно с помощью сканирования.
Реализация такого параллельного алгоритма на основе сканирования намного сложнее (особенно эффективное сканирование, дружественное к кэшу). Относительно простая реализация — дважды выполнить итерацию набора данных, чтобы вычислить локальные суммы, затем вычислить кумулятивные суммы и, наконец, выполнить фактическую работу.
@nb.njit(cache=True, parallel=True)
def numba_parallel_where(arr):
flattenArr = arr.reshape(-1)
arrSize = flattenArr.size
chunkSize = 2048
chunkCount = (arrSize + chunkSize - 1) // chunkSize
counts = np.empty(chunkCount, dtype=np.int32)
for chunkId in nb.prange(chunkCount):
start = chunkId * chunkSize
end = min(start + chunkSize, arrSize)
count = 0
for i in range(start, end):
count += flattenArr[i] < 0
counts[chunkId] = count
offsets = np.empty(chunkCount + 1, dtype=np.int64)
offsets[0] = 0
for chunkId in range(chunkCount):
offsets[chunkId + 1] = offsets[chunkId] + counts[chunkId]
outSize = offsets[-1]
rows = np.empty(outSize, dtype=np.int32)
cols = np.empty(outSize, dtype=np.int32)
n = np.int32(arr.shape[1])
for chunkId in nb.prange(chunkCount):
start = chunkId * chunkSize
end = min(start + chunkSize, arrSize)
idx = offsets[chunkId]
for i in range(start, end):
if flattenArr[i] < 0:
rows[idx] = i // n
cols[idx] = i % n
idx += 1
return rows, cols
Результаты производительности
Вот результаты производительности моего 6-ядерного процессора i5-9600KF с оперативной памятью 40 ГиБ/с (оптимальная пропускная способность для чтения) в Windows:
Sequential: 1.68 ms
Parallel (no pre-allocations): 0.85 ms
Parallel (with pre-allocations): 0.55 ms <-----
Theoretical optimal: 0.20 ms
Параллельная реализация примерно в два раза быстрее. На моей машине удалось достичь пропускной способности памяти ~19 ГиБ/с, что неоптимально, но не так уж и плохо. Хотя на 6-ядерном процессоре это кажется разочаровывающим, параллельная реализация не очень оптимизирована, и есть несколько моментов, которые следует учитывать:
- выделение вывода происходит очень медленно (кажется, занимает 0,2-0,3 мс)
- входной массив читается дважды, оказывая большую нагрузку на память
- модуль/деление, используемые для вычисления индексов, немного дороже
Предварительное выделение выходных данных (заполненных вручную нулями, чтобы избежать возможных ошибок страниц) и передача их в функцию помогает значительно сократить время выполнения. Эта оптимизированная версия в 3 раза быстрее исходной (последовательной). Он достигает ~28 ГиБ/с, что относительно неплохо.
Примечания о производительности
Другие возможные оптимизации включают в себя:
- использование кластеров фрагментов, чтобы избежать двойной итерации всего набора данных (в Numba сложно написать быструю реализацию, поскольку параллелизм очень ограничен по сравнению с родными языками);
- перебирать строки матрицы, чтобы избежать модулей/делений, предполагая, что количество строк не слишком велико и не мало для повышения производительности.
AFAIK, уже несколько лет существует открытая ошибка медленного распределения Numba, но она, очевидно, еще не решена. Кроме того, относительно большие выделения могут быть довольно дорогими и в Windows. Код, скомпилированный в собственном коде (написанный на C/C++/Fortran/Rust) в Linux, не должен вызывать такой проблемы.
Обратите внимание, что такие оптимизации сделают реализацию еще более сложной, чем описанная выше.
Примерно да. значение idx зависит от того, были ли выполнены предыдущие условия или нет. Современные процессоры очень «умны», поскольку они могут предполагать, что условие здесь часто не учитывается (благодаря предсказанию ветвей) на основе внутреннего цикла обратной связи. Исходя из этого, они могут очень быстро протестировать многие состояния arr[i, j] < 0. Некоторые из них можно даже тестировать параллельно на последних процессорах. Однако если одно из условий оказывается неверно предсказанным, это катастрофа для производительности. Действительно, процессору необходимо отменить операции и повторно выполнить код из-за ошибочного прогнозирования.
Эта большая задержка неправильного прогнозирования (например,> 10 циклов) происходит на критическом пути: цепочке зависимостей idx инструкций увеличения. ЦП не может хранить ничего, зависящее от idx, пока оно не исправлено. Когда arr[i, j] < 0 часто бывает ложным (как здесь), процессоры могут почти игнорировать цепочку зависимостей и параллельно выполнять инструкции и конвейеризировать их. Когда это часто так, цепочка зависимостей — это увеличение idx. Когда результат не предсказуем, цепочка зависимостей представляет собой idx условное увеличение + задержку условной ветви.
По этой причине производительность последовательных/параллельных версий зависит от вероятности arr[i, j] < 0 и даже от ее предсказуемости. Параллельная версия должна масштабироваться лучше, когда arr[i, j] < 0 непредсказуемо (потому что здесь больше работы, связанной с ЦП, которая досадно параллельна).
Короче говоря: цепочка зависимостей содержит idx инструкцию увеличения, но мы также можем рассматривать условные ветки как часть цепочки зависимостей в отношении практического использования.
Альтернативой может быть создание пустого массива длиной arr.shape[0]*arr.shape[1], запись в него по индексу i * arr.shape[0] + j, если arr[i, j] < 0, тем самым полностью устранив зависимость? Тогда потребуется еще один проход для перемещения всех записанных значений в начало массива, чтобы это не приводило к ускорению.
Да, вы можете дополнительно оптимизировать последовательный код, обычно используя условное mov, чтобы уменьшить длину критического пути (в терминах тактов). Но в конечном итоге все равно останется длинная цепочка зависимостей. Вы можете сократить его небольшими порциями, как в параллельной версии, но это довольно сложно сделать эффективно. Инструкции SIMD могут помочь и в этом случае. При этом я не думаю, что ваше решение работает. Проблема состоит в том, чтобы собрать результат так, чтобы элементы хранились последовательно. В вашем решении элементы хранятся редко. Это собрание — самая дорогая часть.
Часть перед сбором должна быть распараллеливаема, поскольку я надеялся, что нет никакой зависимости. Но собрание в конце, как вы говорите, является узким местом.
Другой проход — это проверка цикла if arr[i] == 0 и последующее увеличение idx, чтобы переместить значение в result[idx]. Примерно это и делает первоначальный алгоритм ;) .
Часть до сборки не такая уж и дорогая, и современный процессор может выполнять инструкцию arr[i, j] < 0 очень быстро параллельно, поскольку нет никаких зависимостей. Ваше решение, скорее всего, будет медленнее, поскольку для него требуется записать больше данных, а затем снова прочитать их из ОЗУ. Однако это может быть немного быстрее, если вы можете делать это по частям (в кеше L1) и если вы используете инструкции SIMD в первой части. Обратите внимание, что наборы инструкций SIMD, такие как AVX-512, позволяют векторизовать вторую часть.
Последнюю часть сложно распараллелить с использованием нескольких потоков, и решение состоит в том, чтобы использовать метод из этого ответа;).
На моей машине это происходит намного быстрее! %timeit numba_parallel_where(A) 657 мкс ± 35,9 мкс на цикл (среднее ± стандартное отклонение для 7 циклов по 1000 циклов каждый)
Цепочка зависимостей только от idx? То есть только откуда поставить вывод?