cusolver_cholesky_linear_system_solve_potrf_potrs.py

1import numpy as np
2from cuda import cuda, cusolver
3
4# 1. Initialize data on Host (CPU)
5# Matrix A (Positive Definite)
6A = np.array([[1.0, 2.0], [2.0, 5.0]], dtype=np.float64)
7# Right hand side b
8b = np.array([1.0, 1.0], dtype=np.float64)
9n = A.shape[0]
10
11# 2. Initialize CUDA and cuSOLVER
12(err,) = cuda.cuInit(0)
13res, dev = cuda.cuDeviceGet(0)
14res, ctx = cuda.cuCtxCreate(0, dev)
15handle = cusolver.cusolverDnCreate()
16
17# 3. Allocate and copy data to Device (GPU)
18# Allocate A
19res, d_A = cuda.cuMemAlloc(A.nbytes)
20res = cuda.cuMemcpyHtoD(d_A, A.ctypes.data, A.nbytes)
21
22# Allocate b (which will be overwritten by solution x)
23res, d_b = cuda.cuMemAlloc(b.nbytes)
24res = cuda.cuMemcpyHtoD(d_b, b.ctypes.data, b.nbytes)
25
26# 4. Linear System Solve (Cholesky factorization)
27# Get workspace size
28uplo = cusolver.cublasFillMode_t.CUBLAS_FILL_MODE_LOWER
29workspaceSize = cusolver.cusolverDnDpotrf_bufferSize(handle, uplo, n, d_A, n)
30
31# Allocate workspace and info
32res, d_workspace = cuda.cuMemAlloc(workspaceSize)
33res, d_info = cuda.cuMemAlloc(np.int32().nbytes)
34
35# Factorize A = L*L^T
36cusolver.cusolverDnDpotrf(handle, uplo, n, d_A, n, d_workspace, workspaceSize, d_info)
37
38# Solve Ax = b -> L*L^T*x = b
39cusolver.cusolverDnDpotrs(handle, uplo, n, 1, d_A, n, d_b, n, d_info)
40
41# 5. Copy result back to Host
42x = np.empty_like(b)
43cuda.cuMemcpyDtoH(x.ctypes.data, d_b, b.nbytes)
44
45print("Solution x:", x)
46
47# 6. Cleanup
48cusolver.cusolverDnDestroy(handle)
49cuda.cuMemFree(d_A)
50cuda.cuMemFree(d_b)
51cuda.cuMemFree(d_workspace)
52cuda.cuMemFree(d_info)
53cuda.cuCtxDestroy(ctx)