cusparse_sparse_matrix_vector_multiplication_csr_format.py

1import numpy as np
2from cuda import cusparse, cuda
3
4def check_cuda_status(status):
5    if status != cuda.CUresult.CUDA_SUCCESS:
6        raise RuntimeError(f"CUDA Error: {status}")
7
8def check_cusparse_status(status):
9    if status != cusparse.cusparseStatus_t.CUSPARSE_STATUS_SUCCESS:
10        raise RuntimeError(f"cuSPARSE Error: {status}")
11
12# Matrix A (CSR format)
13# [1 0 2]
14# [0 3 0]
15# [4 5 6]
16h_csr_row_ptr = np.array([0, 2, 3, 6], dtype=np.int32)
17h_csr_col_ind = np.array([0, 2, 1, 0, 1, 2], dtype=np.int32)
18h_csr_values  = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], dtype=np.float32)
19h_x = np.array([1.0, 1.0, 1.0], dtype=np.float32)
20h_y = np.array([0.0, 0.0, 0.0], dtype=np.float32)
21
22alpha = 1.0
23beta  = 0.0
24
25# Device memory allocation
26d_csr_row_ptr = cuda.cuMemAlloc(h_csr_row_ptr.nbytes)[1]
27d_csr_col_ind = cuda.cuMemAlloc(h_csr_col_ind.nbytes)[1]
28d_csr_values  = cuda.cuMemAlloc(h_csr_values.nbytes)[1]
29d_x           = cuda.cuMemAlloc(h_x.nbytes)[1]
30d_y           = cuda.cuMemAlloc(h_y.nbytes)[1]
31
32# Copy data to device
33cuda.cuMemcpyHtoD(d_csr_row_ptr, h_csr_row_ptr, h_csr_row_ptr.nbytes)
34cuda.cuMemcpyHtoD(d_csr_col_ind, h_csr_col_ind, h_csr_col_ind.nbytes)
35cuda.cuMemcpyHtoD(d_csr_values, h_csr_values, h_csr_values.nbytes)
36cuda.cuMemcpyHtoD(d_x, h_x, h_x.nbytes)
37
38# Create cuSPARSE handle
39status, handle = cusparse.cusparseCreate()
40
41# Create sparse matrix and dense vectors
42status, matA = cusparse.cusparseCreateCsr(3, 3, 6, d_csr_row_ptr, d_csr_col_ind, d_csr_values,
43                                          cusparse.cusparseIndexType_t.CUSPARSE_INDEX_32I,
44                                          cusparse.cusparseIndexType_t.CUSPARSE_INDEX_32I,
45                                          cusparse.cusparseIndexBase_t.CUSPARSE_INDEX_BASE_ZERO,
46                                          cuda.cudaDataType.CUDA_R_32F)
47status, vecX = cusparse.cusparseCreateDnVec(3, d_x, cuda.cudaDataType.CUDA_R_32F)
48status, vecY = cusparse.cusparseCreateDnVec(3, d_y, cuda.cudaDataType.CUDA_R_32F)
49
50# Buffer size calculation
51status, bufferSize = cusparse.cusparseSpMV_bufferSize(
52    handle, cusparse.cusparseOperation_t.CUSPARSE_OPERATION_NON_TRANSPOSE,
53    np.array([alpha], dtype=np.float32), matA, vecX, np.array([beta], dtype=np.float32), vecY,
54    cuda.cudaDataType.CUDA_R_32F, cusparse.cusparseSpMVAlg_t.CUSPARSE_SPMV_ALG_DEFAULT)
55
56d_buffer = cuda.cuMemAlloc(bufferSize)[1]
57
58# Execution
59status = cusparse.cusparseSpMV(
60    handle, cusparse.cusparseOperation_t.CUSPARSE_OPERATION_NON_TRANSPOSE,
61    np.array([alpha], dtype=np.float32), matA, vecX, np.array([beta], dtype=np.float32), vecY,
62    cuda.cudaDataType.CUDA_R_32F, cusparse.cusparseSpMVAlg_t.CUSPARSE_SPMV_ALG_DEFAULT, d_buffer)
63
64# Copy result back to host
65cuda.cuMemcpyDtoH(h_y, d_y, h_y.nbytes)
66
67print(f"Result y: {h_y}")
68
69# Cleanup
70cusparse.cusparseDestroySpMat(matA)
71cusparse.cusparseDestroyDnVec(vecX)
72cusparse.cusparseDestroyDnVec(vecY)
73cusparse.cusparseDestroy(handle)
74cuda.cuMemFree(d_csr_row_ptr)
75cuda.cuMemFree(d_csr_col_ind)
76cuda.cuMemFree(d_csr_values)
77cuda.cuMemFree(d_x)
78cuda.cuMemFree(d_y)
79cuda.cuMemFree(d_buffer)