在 PyCuda 中就地转置 3D 数组
Inplace transpose of 3D array in PyCuda
我有一个 3D 数组,想转置它的前两个维度 (x & y),但不转置第三个维度 (z)。在 3D 数组 A 上,我想要与 numpy 的 A.transpose((1,0,2))
相同的结果。具体来说,我想获得 "transposed" 全局 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]
...
谢谢,
使用与 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个字节。
我有一个 3D 数组,想转置它的前两个维度 (x & y),但不转置第三个维度 (z)。在 3D 数组 A 上,我想要与 numpy 的 A.transpose((1,0,2))
相同的结果。具体来说,我想获得 "transposed" 全局 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]
...
谢谢,
使用与 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个字节。