Как добиться объединения памяти при переборе неквадратных 2D-массивов?
Я изо всех сил пытаюсь понять слияние памяти в CUDA. Чтобы оценить разницу в производительности между объединенным и несвязанным доступом к памяти, я реализовал две разные версии ядра, которые добавляют две 2D-матрицы:
from numba import cuda
@cuda.jit
def uncoalesced_matrix_add(a, b, out):
x, y = cuda.grid(2)
out[x][y] = a[x][y] + b[x][y]
@cuda.jit
def coalesced_matrix_add(a, b, out):
x, y = cuda.grid(2)
out[y][x] = a[y][x] + b[y][x]
Когда я тестирую приведенный выше код с квадратными матрицами, все работает нормально, я имею в виду, что оба ядра дают одинаковый результат, а объединенная версия работает значительно быстрее:
import numpy as np
nrows, ncols = 2048, 2048
tpb = 32
threads_per_block = (tpb, tpb)
blocks = ((nrows + (tpb - 1))//tpb, (ncols + (tpb - 1))//tpb)
size = nrows*ncols
a = np.arange(size).reshape(nrows, ncols).astype(np.int32)
b = np.ones(shape=a.shape, dtype=np.int32)
out = np.empty_like(a).astype(np.int32)
d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_out = cuda.to_device(out)
uncoalesced_matrix_add[blocks, threads_per_block](d_a, d_b, d_out)
slow = d_out.copy_to_host()
coalesced_matrix_add[blocks, threads_per_block](d_a, d_b, d_out)
fast = d_out.copy_to_host()
np.array_equal(slow, fast)
# True
Однако, если я изменю nrows = 1024 так, чтобы матрицы больше не были квадратными, coalesced_matrix_add() выдаст следующую ошибку:
CudaAPIError: [700] Call to cuMemcpyDtoH results in UNKNOWN_CUDA_ERROR
Что мне здесь не хватает?
Редактировать
Для полноты прилагаю профилирование. Эти данные были получены с помощью обходного пути, предложенного @Robert Crovella с nrows = 1024 и ncols = 2048:
In [40]: %timeit uncoalesced_matrix_add[blocksu, threads_per_block](d_a, d_b, d_out)
289 µs ± 498 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [41]: %timeit coalesced_matrix_add[blocksc, threads_per_block](d_a, d_b, d_out)
164 µs ± 108 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Ответы 1
когда вы делаете свой массив неквадратным, этот расчет больше не является правильным для один ваших двух ядер:
blocks = ((nrows + (tpb - 1))//tpb, (ncols + (tpb - 1))//tpb)
Как для потоков, так и для блоков размерность индекса x идет первой, за которой следует размерность индекса y (ссылаясь на x и y, поскольку они появляются во встроенных в ядро переменных x и y).
Для использования в вашем первом ядре:
out[x][y] = a[x][y] + b[x][y]
мы хотим, чтобы x индексировал строки. Это согласуется с вашим определением сетки.
Для использования во втором ядре:
out[y][x] = a[y][x] + b[y][x]
мы хотим, чтобы y индексировал строки. Это не соответствует вашему определению сетки.
Результатом является доступ за пределами границ при втором вызове ядра. Ориентация вашей прямоугольной сетки не соответствует ориентации ваших прямоугольных данных.
В случае квадрата такое обращение было несущественным, так как оба измерения были одинаковыми.
Вот возможное "исправление":
$ cat t62.py
from numba import cuda
import numpy as np
@cuda.jit
def uncoalesced_matrix_add(a, b, out):
x, y = cuda.grid(2)
out[x][y] = a[x][y] + b[x][y]
@cuda.jit
def coalesced_matrix_add(a, b, out):
x, y = cuda.grid(2)
out[y][x] = a[y][x] + b[y][x]
nrows, ncols = 512, 1024
tpb = 32
threads_per_block = (tpb, tpb)
blocksu = ((nrows + (tpb - 1))//tpb, (ncols + (tpb - 1))//tpb)
blocksc = ((ncols + (tpb - 1))//tpb, (nrows + (tpb - 1))//tpb)
size = nrows*ncols
a = np.arange(size).reshape(nrows, ncols).astype(np.int32)
b = np.ones(shape=a.shape, dtype=np.int32)
out = np.empty_like(a).astype(np.int32)
d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_out = cuda.to_device(out)
uncoalesced_matrix_add[blocksu, threads_per_block](d_a, d_b, d_out)
slow = d_out.copy_to_host()
coalesced_matrix_add[blocksc, threads_per_block](d_a, d_b, d_out)
fast = d_out.copy_to_host()
print(np.array_equal(slow, fast))
# True
$ python t62.py
True
$
Также обратите внимание, что эта стратегия определения размера сетки работает только для измерений, целое число которых делится на размер блока.
Ваше решение сработало как шарм. Хотел бы я проголосовать больше одного раза :)