在 PyCuda 中就地转置 3D 数组答案

【问题标题】：Inplace transpose of 3D array in PyCuda在 PyCuda 中就地转置 3D 数组
【发布时间】：2015-11-14 15:27:21
【问题描述】：

我有一个 3D 数组，想转置它的前两个维度 (x & y)，而不是第三个 (z)。在 3D 数组 A 上，我想要与 numpy 的 A.transpose((1,0,2)) 相同的结果。具体来说，我想获得“转置”的全局threadIdx。下面的代码应该在 3D 数组 A 中的未转置位置写入转置索引。它没有。

有什么建议吗？

import numpy as np
from pycuda import compiler, gpuarray
import pycuda.driver as cuda
import pycuda.autoinit

kernel_code = """
__global__ void test_indexTranspose(uint*A){
    const size_t size_x = 4;
    const size_t size_y = 4;
    const size_t size_z = 3;

    // Thread position in each dimension
    const size_t tx = blockDim.x * blockIdx.x + threadIdx.x;
    const size_t ty = blockDim.y * blockIdx.y + threadIdx.y;
    const size_t tz = blockDim.z * blockIdx.z + threadIdx.z;

    if(tx < size_x && ty < size_y && tz < size_z){
        // Flat index
        const size_t ti = tz * size_x * size_y + ty * size_x + tx;
        // Transposed flat index
        const size_t tiT = tz * size_x * size_y + tx * size_x + ty;
        A[ti] = tiT;
    }
}
"""

A = np.zeros((4,4,3),dtype=np.uint32)
mod = compiler.SourceModule(kernel_code)
test_indexTranspose = mod.get_function('test_indexTranspose')
A_gpu = gpuarray.to_gpu(A)
test_indexTranspose(A_gpu, block=(2, 2, 1), grid=(2,2,3))

这是返回的（不是我期望的）：

A_gpu.get()[:,:,0]
array([[ 0, 12,  9,  6],
       [ 3, 15, 24, 21],
       [18, 30, 27, 36],
       [33, 45, 42, 39]], dtype=uint32)

A_gpu.get()[:,:,1]
array([[ 4,  1, 13, 10],
       [ 7, 16, 28, 25],
       [22, 19, 31, 40],
       [37, 34, 46, 43]], dtype=uint32)

A_gpu.get()[:,:,2]
array([[ 8,  5,  2, 14],
       [11, 20, 17, 29],
       [26, 23, 32, 44],
       [41, 38, 35, 47]], dtype=uint32)

这是我所期望的（但没有返回）：

A_gpu.get()[:,:,0]
array([[0, 4, 8,  12],
       [1, 5, 9,  13],
       [2, 6, 10, 14],
       [3, 7, 11, 15]], dtype=uint32)

A_gpu.get()[:,:,1]
array([[16, 20, 24, 28],
       [17, 21, 25, 29],
       [18, 22, 26, 30],
       [19, 23, 27, 31]], dtype=uint32)

A_gpu.get()[:,:,2]
...

谢谢，

【问题讨论】：

这似乎是由于 numpy.ndarray 的跨步/内存布局。该代码假定 x 变化最快，然后 y 和 z 最慢。但是，A.strides 是 12, 3, 1 (y, x, z)，根据 numpy.ndarray 文档，C 顺序和默认值，表明 z 变化最快。一种可能的解决方案是创建具有所需步幅的数组。即：A = np.ndarray(shape=(4,4,3),dtype=np.uint32,strides=(16,4,64))。但是，这是否有效，我明天回到实验室计算机时就知道了。

标签： python numpy multidimensional-array cuda pycuda

【解决方案1】：

创建具有与 CUDA 内核代码一致的步幅的 numpy 数组可以解决问题。 numpy 数组的默认布局不是我的内核假设的行、列、深度。但是，可以在创建数组时设置步幅。
如果数组是这样创建的，上面的内核可以正常工作：

nRows = 4
nCols = 4
nSlices = 3
nBytes = np.dtype(np.uint32).itemsize
A = np.ndarray(shape=(nRows, nCols, nSlices), 
               dtype=np.uint32, 
               strides=(nCols*nBytes, 1*nBytes, nCols*nRows*nBytes))

步幅是每个维度的连续索引在内存中的跳跃（以字节为单位）。例如。从第 1 行的第一个元素到第 2 行的第一个元素有nCols * nBytes，即 16 个字节。

【讨论】：