Несовпадение в двумерной функции CUDA-FFTShift ⇐ Python

[Anonymous]

Я новичок в программировании CUDA, и мне нужно выполнить операцию сдвига fft над плоским двумерным массивом. Я немного поискал и наткнулся на эту библиотеку, но до сих пор не смог заставить ее работать даже после многочисленных попыток. Вывод:

• Как-то неправильно выровнен.
• "Поврежден", вероятно, из-за того, что потоки мешают друг другу.

Я решил написать локальную 2D-функцию ffshift на Python, и там она работает отлично. Я не могу понять, что я делаю не так в версии CUDA. В целом я не эксперт по функции fftshift, но тот факт, что эта версия работает на Python, меня сбивает с толку.

Код: Выделить всё

template 
__global__ void cufftShift_2D_kernel(T* array, int N)
{
int x = blockIdx.x * blockDim.x + threadIdx.x;
int y = blockIdx.y * blockDim.y + threadIdx.y;

int index = y * N + x;

int offsetA = (N * N + N) / 2;
int offsetB = (N * N - N) / 2;

T temp;

if (x < N / 2) {
if (y < N / 2) {
temp = array[index];
array[index] = array[index + offsetA];
array[index + offsetA] = temp;
}
}
else if (y < N / 2) {
temp = array[index];
array[index] = array[index + offsetB];
array[index + offsetB] = temp;
}
}

template 
struct GPUBuffer {
thrust::device_vector buffer;
T* p;
std::size_t mem_size;

GPUBuffer() = delete;

GPUBuffer(std::size_t size) :
buffer(thrust::device_vector(size)),
p(thrust::raw_pointer_cast(buffer.data())),
mem_size(buffer.size() * sizeof(T)) {}
};

__global__ void multiplyBuffers(cufftComplex *kernel, cufftComplex *field, int size) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;

if (idx < size) {
cufftComplex x = kernel[idx];
cufftComplex y = field[idx];

field[idx].x = x.x * y.x - x.y * y.y;
field[idx].y = x.x * y.y + x.y * y.x;
}
}

int main()
{
const int field_size = 32 * scale;

std::vector field = std::vector(field_size * field_size, 0.f);

const int threadsPerBlock = 256;
const int blocksPerGrid = (field.size() + threadsPerBlock - 1) / threadsPerBlock;

cufftHandle normal, inv;

cufftPlan2d(&normal, field_size, field_size, CUFFT_C2C);
cufftPlan2d(&inv, field_size, field_size, CUFFT_C2C);

GPUBuffer kernel_gpu(field.size());
GPUBuffer field_gpu(field.size());
GPUBuffer shift_gpu(field.size());

std::vector host_output(field.size());

std::vector real_output(field.size());

for (size_t i = 0; i < field.size(); i++)
{
field_gpu.buffer[i] = { static_cast(i), 0.f };
kernel_gpu.buffer[i] = { static_cast(i), 0.f };
}

cufftExecC2C(normal, kernel_gpu.p, kernel_gpu.p, CUFFT_FORWARD);

for (int i = 0; i < 1; ++i) {
cufftExecC2C(normal, field_gpu.p, field_gpu.p, CUFFT_FORWARD);
multiplyBuffers(kernel_gpu.p, field_gpu.p, field.size());
cufftExecC2C(inv, field_gpu.p, field_gpu.p, CUFFT_INVERSE);
cudaMemcpy(shift_gpu.p, field_gpu.p, field_gpu.mem_size, cudaMemcpyDeviceToDevice);
cufftShift_2D_kernel(shift_gpu.p, field_size);
}

cudaMemcpy(host_output.data(), field_gpu.p, field_gpu.mem_size, cudaMemcpyDeviceToHost);

for (size_t i = 0; i < field.size(); ++i) {
real_output[i] = sqrt(host_output[i].x * host_output[i].x + host_output[i].y * host_output[i].y);
}

dump_array_to_file(real_output, field_size, field_size, "cuda_inv.txt");

cudaMemcpy(host_output.data(), shift_gpu.p, shift_gpu.mem_size, cudaMemcpyDeviceToHost);

for (size_t i = 0; i < field.size(); ++i) {
real_output[i] = sqrt(host_output[i].x * host_output[i].x + host_output[i].y * host_output[i].y);
}

dump_array_to_file(real_output, field_size, field_size, "cuda_shifted.txt");

cufftDestroy(normal);
cufftDestroy(inv);
return 0;
}

Вот эквивалентная версия Python, которая выдает правильный результат:

Код: Выделить всё

def shift(array: NDArray):
N = np.int32(np.sqrt(array.size))
offsetA = (N * N + N) // 2
offsetB = (N * N - N) // 2
for (y, x), _ in np.ndenumerate(array.reshape((N, N))):
idx = y * N + x
if x < N // 2:
if y < N // 2:
array[idx], array[idx + offsetA] = array[idx + offsetA], array[idx]
else:
if y <  N // 2:
array[idx], array[idx + offsetB] = array[idx + offsetB], array[idx]
return array.reshape((N, N))

И как это выглядит:




Подробнее здесь: https://stackoverflow.com/questions/792 ... t-function