nvidia_curand_cuda_python_random_uniform_float_generation.py

1import numpy as np
2from cuda import cuda, curand
3
4def check_cuda_error(err):
5    if isinstance(err, cuda.CUresult):
6        if err != cuda.CUresult.CUDA_SUCCESS:
7            raise RuntimeError(f"CUDA Error: {err}")
8    elif isinstance(err, curand.curandStatus):
9        if err != curand.curandStatus.CURAND_STATUS_SUCCESS:
10            raise RuntimeError(f"cuRAND Error: {err}")
11
12# 1. Initialize CUDA Driver API
13err, = cuda.cuInit(0)
14check_cuda_error(err)
15
16err, device = cuda.cuDeviceGet(0)
17check_cuda_error(err)
18
19err, context = cuda.cuCtxCreate(0, device)
20check_cuda_error(err)
21
22# 2. Parameters
23n = 10
24host_data = np.zeros(n, dtype=np.float32)
25
26# 3. Allocate device memory
27err, device_data = cuda.cuMemAlloc(host_data.nbytes)
28check_cuda_error(err)
29
30# 4. Create cuRAND generator
31err, generator = curand.curandCreateGenerator(curand.curandRngType.CURAND_RNG_PSEUDO_DEFAULT)
32check_cuda_error(err)
33
34# 5. Set seed
35err = curand.curandSetPseudoRandomGeneratorSeed(generator, 1234)
36check_cuda_error(err)
37
38# 6. Generate random numbers on device
39err = curand.curandGenerateUniform(generator, device_data, n)
40check_cuda_error(err)
41
42# 7. Copy results back to host
43err = cuda.cuMemcpyDtoH(host_data, device_data, host_data.nbytes)
44check_cuda_error(err)
45
46print("Generated Random Numbers:")
47print(host_data)
48
49# 8. Cleanup
50curand.curandDestroyGenerator(generator)
51cuda.cuMemFree(device_data)
52cuda.cuCtxDestroy(context)