Для этого вам нужно описать расположение массива в памяти для ядра CUDA, и вам нужны правильные вычисления индексации в ядре с использованием предоставленных шагов на стороне хоста. Простой способ сделать это - определить небольшой вспомогательный класс в CUDA, который скрывает основную часть индексации и обеспечивает простой синтаксис индексации. Например:
from pycuda import driver, gpuarray
from pycuda.compiler import SourceModule
import pycuda.autoinit
import numpy as np
mod=SourceModule("""
struct stride3D
{
float* p;
int s0, s1;
__device__
stride3D(float* _p, int _s0, int _s1) : p(_p), s0(_s0), s1(_s1) {};
__device__
float operator () (int x, int y, int z) const { return p[x*s0 + y*s1 + z]; };
__device__
float& operator () (int x, int y, int z) { return p[x*s0 + y*s1 + z]; };
};
__global__ void mat_ops(float *A, int sA0, int sA1, float *B, int sB0, int sB1)
{
stride3D A3D(A, sA0, sA1);
stride3D B3D(B, sB0, sB1);
int xidx = blockIdx.x;
int yidx = threadIdx.x;
int zidx = threadIdx.y;
B3D(xidx, yidx, zidx) = A3D(xidx, yidx, zidx);
}
""")
A = 1 + np.arange(0, 4*4*3, dtype=np.float32).reshape(4,4,3)
B = np.zeros((5,5,5), dtype=np.float32)
A_k = gpuarray.to_gpu(A)
B_k = gpuarray.to_gpu(B)
astrides = np.array(A.strides, dtype=np.int32) // A.itemsize
bstrides = np.array(B.strides, dtype=np.int32) // B.itemsize
func = mod.get_function("mat_ops")
func(A_k, astrides[0], astrides[1], B_k, bstrides[0], bstrides[1], grid=(4,1,1),block=(4,3,1))
print(B_k[:4,:4,:3])
Здесь я решил сделать массивы источника и назначения разными размерами, просто чтобы показать, что код является общим и будет работать для массивов любого размера, пока достаточен размер блока. Обратите внимание, что здесь нет проверки границ массива на стороне кода устройства, вам нужно будет добавить это для нетривиальных примеров.
Обратите внимание также, что это должно работать правильно как для fortran, так и для C упорядоченные массивы numpy, поскольку он напрямую использует значения шага numpy. Однако производительность будет зависеть от CUDA из-за проблем с объединением памяти.
Примечание: это не сработает как для заказов fortran, так и для C без расширения вспомогательного класса, чтобы сделать шаги для всех размеры и изменение ядра для принятия шагов для всех измерений входных и выходных массивов. С точки зрения производительности было бы лучше написать отдельные вспомогательные классы для Fortran и C упорядоченных массивов.