nvidia_cufft_1d_complex_to_complex_fft_python_bindings.py

1import numpy as np
2from cuda import cuda, cufft
3
4def checkCudaStatus(status):
5    if status != cuda.CUresult.CUDA_SUCCESS:
6        raise RuntimeError(f"CUDA Error: {status}")
7
8def checkCufftStatus(status):
9    if status != cufft.cufftResult.CUFFT_SUCCESS:
10        raise RuntimeError(f"cuFFT Error: {status}")
11
12# 1. Initialize data
13N = 1024
14host_data = np.arange(N, dtype=np.complex64)
15device_data_size = host_data.nbytes
16
17# 2. Initialize CUDA and allocate memory
18checkCudaStatus(cuda.cuInit(0))
19device = cuda.cuDeviceGet(0)[1]
20context = cuda.cuCtxCreate(0, device)[1]
21
22err, device_data = cuda.cuMemAlloc(device_data_size)
23checkCudaStatus(err)
24
25# 3. Copy data to device
26checkCudaStatus(cuda.cuMemcpyHtoD(device_data, host_data, device_data_size))
27
28# 4. Create cuFFT plan and execute
29# CUFFT_C2C represents Complex-to-Complex transform
30plan = cufft.cufftHandle()
31checkCufftStatus(cufft.cufftPlan1d(plan, N, cufft.cufftType.CUFFT_C2C, 1))
32
33# Execute forward transform
34checkCufftStatus(cufft.cufftExecC2C(plan, device_data, device_data, cufft.CUFFT_FORWARD))
35
36# 5. Copy results back to host
37output_data = np.empty_like(host_data)
38checkCudaStatus(cuda.cuMemcpyDtoH(output_data, device_data, device_data_size))
39
40# 6. Cleanup
41checkCufftStatus(cufft.cufftDestroy(plan))
42checkCudaStatus(cuda.cuMemFree(device_data))
43checkCudaStatus(cuda.cuCtxDestroy(context))
44
45print("Successfully performed FFT.")
46print(f"First 5 elements of output: {output_data[:5]}")