nvidia_cufft_1d_forward_fft_complex_transform.py

1import numpy as np
2from cuda import cuda, cufft
3
4# 1. Initialize data
5n = 1024
6host_data = np.arange(n, dtype=np.complex64)
7device_data_size = host_data.nbytes
8
9# 2. Allocate device memory
10err, = cuda.cuInit(0)
11err, device = cuda.cuDeviceGet(0)
12err, context = cuda.cuCtxCreate(0, device)
13err, device_data = cuda.cuMemAlloc(device_data_size)
14
15# 3. Copy data to device
16err, = cuda.cuMemcpyHtoD(device_data, host_data.ctypes.data, device_data_size)
17
18# 4. Create cuFFT plan and execute
19# CUFFT_C2C represents Complex-to-Complex transform
20err, plan = cufft.cufftPlan1d(n, cufft.cufftType.CUFFT_C2C, 1)
21err = cufft.cufftExecC2C(plan, device_data, device_data, cufft.CUFFT_FORWARD)
22
23# 5. Copy results back to host
24host_result = np.empty_like(host_data)
25err, = cuda.cuMemcpyDtoH(host_result.ctypes.data, device_data, device_data_size)
26
27# 6. Cleanup
28err = cufft.cufftDestroy(plan)
29err = cuda.cuMemFree(device_data)
30err = cuda.cuCtxDestroy(context)
31
32print("FFT Result (first 5 elements):")
33print(host_result[:5])