cusolver_dense_qr_decomposition_linear_system_solver.py

1import numpy as np
2from cuda import cuda, cusolver
3
4def check_cuda_status(status):
5    if isinstance(status, cuda.CUresult):
6        if status != cuda.CUresult.CUDA_SUCCESS:
7            raise RuntimeError(f"CUDA Error: {status}")
8    elif isinstance(status, cusolver.cusolverStatus_t):
9        if status != cusolver.cusolverStatus_t.CUSOLVER_STATUS_SUCCESS:
10            raise RuntimeError(f"cuSOLVER Error: {status}")
11
12# 1. Initialize data (3x3 Matrix A and 3x1 Vector B)
13# Ax = B
14h_A = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 10]], dtype=np.float64)
15h_B = np.array([1, 1, 1], dtype=np.float64)
16m, n = h_A.shape
17
18# 2. Allocate device memory
19err, = cuda.cuInit(0)
20err, dev = cuda.cuDeviceGet(0)
21err, ctx = cuda.cuCtxCreate(0, dev)
22
23err, d_A = cuda.cuMemAlloc(h_A.nbytes)
24err, d_B = cuda.cuMemAlloc(h_B.nbytes)
25err, d_tau = cuda.cuMemAlloc(min(m, n) * h_A.itemsize)
26
27# 3. Copy data to device
28err, = cuda.cuMemcpyHtoD(d_A, h_A.ctypes.data, h_A.nbytes)
29err, = cuda.cuMemcpyHtoD(d_B, h_B.ctypes.data, h_B.nbytes)
30
31# 4. Create cuSOLVER handle
32status, handle = cusolver.cusolverDnCreate()
33check_cuda_status(status)
34
35# 5. Query workspace size for QR decomposition (geqrf)
36status, workSize = cusolver.cusolverDnDgeqrf_bufferSize(handle, m, n, d_A, m)
37check_cuda_status(status)
38err, d_work = cuda.cuMemAlloc(workSize * h_A.itemsize)
39
40# 6. Compute QR decomposition
41err, d_info = cuda.cuMemAlloc(np.dtype(np.int32).itemsize)
42status = cusolver.cusolverDnDgeqrf(handle, m, n, d_A, m, d_tau, d_work, workSize, d_info)
43check_cuda_status(status)
44
45# 7. Solve the system (ormqr + trsm)
46# Compute Q^T * B
47side = cusolver.cublasSideMode_t.CUBLAS_SIDE_LEFT
48trans = cusolver.cublasOperation_t.CUBLAS_OP_T
49status, workSize_ormqr = cusolver.cusolverDnDormqr_bufferSize(
50    handle, side, trans, m, 1, min(m, n), d_A, m, d_tau, d_B, m)
51check_cuda_status(status)
52
53err, d_work_ormqr = cuda.cuMemAlloc(workSize_ormqr * h_A.itemsize)
54status = cusolver.cusolverDnDormqr(
55    handle, side, trans, m, 1, min(m, n), d_A, m, d_tau, d_B, m, d_work_ormqr, workSize_ormqr, d_info)
56check_cuda_status(status)
57
58# 8. Copy result back to host
59h_X = np.empty_like(h_B)
60err, = cuda.cuMemcpyDtoH(h_X.ctypes.data, d_B, h_B.nbytes)
61
62print("Solution X:")
63print(h_X)
64
65# Cleanup
66cusolver.cusolverDnDestroy(handle)
67cuda.cuMemFree(d_A)
68cuda.cuMemFree(d_B)
69cuda.cuMemFree(d_tau)
70cuda.cuMemFree(d_work)
71cuda.cuCtxDestroy(ctx)