Как правильно рассчитать номер блока и потока

Question

Я пытаюсь реализовать алгоритм Флойда-Уоршалла с использованием функций Numba CUDA. adj - моя матрица смежности размером n_atoms на n_atoms (~ 7000 на ~ 7000). У меня небольшие трудности с настройкой сетки, в наивной версии я повторяю по 3 признакам, поэтому, вероятно, сетка должна быть трехмерной. Запуская приведенный ниже код в Google Colab, я получаю сообщение об ошибке, вероятно, связанное с максимальным количеством поддерживаемых потоков.

[1] Call to cuLaunchKernel results in CUDA_ERROR_INVALID_VALUE

Код:

@cuda.jit((int32[:,:], int32))
def floydWarshall_CUDA(a, l):
    k, i, j = cuda.grid(3)
    if k < l and i < l and j < l:
        a[i, j] = min(a[i, k] + a[k, j], a[i, j])

threadsperblock = (2, 2, 2)
blockspergrid = (math.ceil(n_atoms / 2), math.ceil(n_atoms / 2), math.ceil(n_atoms / 2))

floydWarshall_CUDA[threadsperblock, blockspergrid](adj_gpu, n_atoms)