curand_gpu_uniform_float_generation_with_cuda_driver_api.py

1import numpy as np
2from cuda import cuda, curand
3
4def check_status(status):
5    if status != curand.curandStatus.CURAND_STATUS_SUCCESS:
6        raise RuntimeError(f"cuRAND Error: {status}")
7
8def check_cuda_status(status):
9    if status != cuda.CUresult.CUDA_SUCCESS:
10        raise RuntimeError(f"CUDA Error: {status}")
11
12# 1. Initialize CUDA Driver API
13check_cuda_status(cuda.cuInit(0))
14device_id = 0
15res, device = cuda.cuDeviceGet(device_id)
16res, context = cuda.cuCtxCreate(0, device)
17
18# 2. Parameters
19n = 10
20size = n * np.dtype(np.float32).itemsize
21
22# 3. Allocate memory on GPU
23res, d_data = cuda.cuMemAlloc(size)
24
25# 4. Create cuRAND generator
26res, gen = curand.curandCreateGenerator(curand.curandRngType.CURAND_RNG_PSEUDO_DEFAULT)
27check_status(res)
28
29# 5. Set seed
30check_status(curand.curandSetPseudoRandomGeneratorSeed(gen, 1234))
31
32# 6. Generate n floats on device
33# Note: d_data is cast to int for the C-binding pointer requirement
34check_status(curand.curandGenerateUniform(gen, int(d_data), n))
35
36# 7. Copy results back to host
37host_data = np.zeros(n, dtype=np.float32)
38check_cuda_status(cuda.cuMemcpyDtoH(host_data, d_data, size))
39
40print("Generated Random Numbers:")
41print(host_data)
42
43# 8. Cleanup
44curand.curandDestroyGenerator(gen)
45cuda.cuMemFree(d_data)
46cuda.cuCtxDestroy(context)