CuPy-based conjugate gradient solver for linear systems.
| 40 | |
| 41 | |
| 42 | class CuPyCGSolver: |
| 43 | """CuPy-based conjugate gradient solver for linear systems.""" |
| 44 | |
| 45 | def __init__(self, rtol=1e-5, atol=1e-12, maxiter=None): |
| 46 | self.rtol = rtol |
| 47 | self.atol = atol |
| 48 | self.maxiter = maxiter |
| 49 | self.A = None # Cache the sparse matrix |
| 50 | |
| 51 | def solve(self, **kwargs): |
| 52 | # Extract buffers from kwargs |
| 53 | index_ptr: memoryview = kwargs['index_ptr'] # int32, read-only memoryview |
| 54 | indices: memoryview = kwargs['indices'] # int32, read-only memoryview |
| 55 | values: memoryview = kwargs['values'] # float64, read-only memoryview |
| 56 | rhs: memoryview = kwargs['rhs'] # float64, read-only memoryview |
| 57 | x: memoryview = kwargs['x'] # float64, writeable memoryview |
| 58 | num_eqn: int = kwargs['num_eqn'] |
| 59 | nnz: int = kwargs['nnz'] |
| 60 | matrix_status: str = kwargs['matrix_status'] # UNCHANGED, STRUCTURE_CHANGED, COEFFICIENTS_CHANGED |
| 61 | |
| 62 | # Wrap memoryviews using zero-copy numpy views |
| 63 | indptr = np.frombuffer(index_ptr, dtype=np.int32, count=num_eqn + 1) |
| 64 | idx = np.frombuffer(indices, dtype=np.int32, count=nnz) |
| 65 | vals = np.frombuffer(values, dtype=np.float64, count=nnz) |
| 66 | |
| 67 | # Rebuild matrix if structure changed, update values if coefficients changed |
| 68 | if matrix_status == 'STRUCTURE_CHANGED' or self.A is None: |
| 69 | # Copy the entire CSR matrix to the GPU |
| 70 | values_gpu = cp.asarray(vals) |
| 71 | indices_gpu = cp.asarray(idx) |
| 72 | index_ptr_gpu = cp.asarray(indptr) |
| 73 | self.A = cp.sparse.csr_matrix((values_gpu, indices_gpu, index_ptr_gpu), shape=(num_eqn, num_eqn)) |
| 74 | elif matrix_status == 'COEFFICIENTS_CHANGED': |
| 75 | # Update the values of the CSR matrix on the GPU |
| 76 | values_gpu = cp.asarray(vals) |
| 77 | self.A.data[:] = values_gpu # in-place update |
| 78 | else: |
| 79 | # If UNCHANGED, do nothing |
| 80 | pass |
| 81 | |
| 82 | # Wrap RHS for solving |
| 83 | rhs_buf = np.frombuffer(rhs, dtype=np.float64, count=num_eqn) |
| 84 | rhs_gpu = cp.asarray(rhs_buf) |
| 85 | |
| 86 | # Solve using conjugate gradient (without preconditioning) on the GPU |
| 87 | x_gpu, info = cupyx.scipy.sparse.linalg.cg(self.A, rhs_gpu, tol=self.rtol, atol=self.atol, maxiter=self.maxiter) |
| 88 | |
| 89 | # Copy result back to CPU buffer |
| 90 | x_buf = np.frombuffer(x, dtype=np.float64, count=num_eqn) |
| 91 | x_buf[:] = cp.asnumpy(x_gpu) # in-place update |
| 92 | return -int(info) # Return the info from the solver |
| 93 | |
| 94 | |
| 95 | # ----------------------------------------------------------------------------- |