如何在 Numba 中逐行遍历 3D 矩阵？答案

【问题标题】：How can I traverse a 3D matrix on a thread per row basis in Numba?如何在 Numba 中逐行遍历 3D 矩阵？
【发布时间】：2020-12-02 19:25:30
【问题描述】：

我正在尝试使用 Numba CUDA 实现基于运行的内核，其中我需要在每个线程的行上遍历 3D 矩阵的元素，即每个线程都被分配一个迭代该行的所有元素的行。

例如，为了简单起见，如果我要使用 50 行 100 列的 2D 矩阵，我需要创建 50 个线程来遍历它们各自行的 100 个元素。

谁能告诉我怎么做？

【问题讨论】：

标签： python multidimensional-array gpgpu numba connected-components

【解决方案1】：

事实证明它实际上很简单。您只需要启动与行一样多的线程并让内核“指出”它的方向。这是一个简单的内核，演示了如何在 3D 矩阵 (binary_image) 上进行这样的迭代。内核本身是我正在实现的 CCL 算法的一部分，但可以放心地忽略它：

from numba import cuda

@cuda.jit
def kernel_1(binary_image, image_width, s_matrix, labels_matrix):
    # notice how we're only getting the row and depth of each thread
    row, image_slice = cuda.grid(2) 
    sm_pos, lm_pos = 0, 0
    span_found = False
    if row < binary_image.shape[0] and image_slice < binary_image.shape[2]:  # guard for rows and slices
        # and here's the traversing over the columns
        for column in range(binary_image.shape[1]):
            if binary_image[row, column, image_slice] == 0:
                if not span_found:  # Connected Component found
                    span_found = True
                    s_matrix[row, sm_pos, image_slice] = column
                    sm_pos = sm_pos + 1
                    # converting 2D coordinate to 1D
                    linearized_index = row * image_width + column
                    labels_matrix[row, lm_pos, image_slice] = linearized_index
                    lm_pos = lm_pos + 1
                else:
                    s_matrix[row, sm_pos, image_slice] = column
            elif binary_image[row, column, image_slice] == 255 and span_found:
                span_found = False
                s_matrix[row, sm_pos, image_slice] = column - 1
                sm_pos = sm_pos + 1

【讨论】：