def csrgeam(a, b, alpha=1, beta=1): """Matrix-matrix addition. .. math:: C = \\alpha A + \\beta B Args: a (cupyx.scipy.sparse.csr_matrix): Sparse matrix A. b (cupyx.scipy.sparse.csr_matrix): Sparse matrix B. alpha (float): Coefficient for A. beta (float): Coefficient for B. Returns: cupyx.scipy.sparse.csr_matrix: Result matrix. """ if not check_availability('csrgeam'): raise RuntimeError('csrgeam is not available.') if not isinstance(a, cupyx.scipy.sparse.csr_matrix): raise TypeError('unsupported type (actual: {})'.format(type(a))) if not isinstance(b, cupyx.scipy.sparse.csr_matrix): raise TypeError('unsupported type (actual: {})'.format(type(b))) assert a.has_canonical_format assert b.has_canonical_format if a.shape != b.shape: raise ValueError('inconsistent shapes') handle = device.get_cusparse_handle() m, n = a.shape a, b = _cast_common_type(a, b) nnz = numpy.empty((), 'i') cusparse.setPointerMode(handle, cusparse.CUSPARSE_POINTER_MODE_HOST) c_descr = MatDescriptor.create() c_indptr = cupy.empty(m + 1, 'i') cusparse.xcsrgeamNnz(handle, m, n, a._descr.descriptor, a.nnz, a.indptr.data.ptr, a.indices.data.ptr, b._descr.descriptor, b.nnz, b.indptr.data.ptr, b.indices.data.ptr, c_descr.descriptor, c_indptr.data.ptr, nnz.ctypes.data) c_indices = cupy.empty(int(nnz), 'i') c_data = cupy.empty(int(nnz), a.dtype) alpha = numpy.array(alpha, a.dtype).ctypes beta = numpy.array(beta, a.dtype).ctypes _call_cusparse('csrgeam', a.dtype, handle, m, n, alpha.data, a._descr.descriptor, a.nnz, a.data.data.ptr, a.indptr.data.ptr, a.indices.data.ptr, beta.data, b._descr.descriptor, b.nnz, b.data.data.ptr, b.indptr.data.ptr, b.indices.data.ptr, c_descr.descriptor, c_data.data.ptr, c_indptr.data.ptr, c_indices.data.ptr) c = cupyx.scipy.sparse.csr_matrix((c_data, c_indices, c_indptr), shape=a.shape) c._has_canonical_format = True return c
def csrgemm2(a, b, d=None, alpha=1, beta=1): """Matrix-matrix product for CSR-matrix. math:: C = alpha * A * B + beta * D Args: a (cupyx.scipy.sparse.csr_matrix): Sparse matrix A. b (cupyx.scipy.sparse.csr_matrix): Sparse matrix B. d (cupyx.scipy.sparse.csr_matrix or None): Sparse matrix D. alpha (scalar): Coefficient beta (scalar): Coefficient Returns: cupyx.scipy.sparse.csr_matrix """ if not check_availability('csrgemm2'): raise RuntimeError('csrgemm2 is not available.') assert a.ndim == b.ndim == 2 if not isinstance(a, cupyx.scipy.sparse.csr_matrix): raise TypeError('unsupported type (actual: {})'.format(type(a))) if not isinstance(b, cupyx.scipy.sparse.csr_matrix): raise TypeError('unsupported type (actual: {})'.format(type(b))) assert a.has_canonical_format assert b.has_canonical_format if a.shape[1] != b.shape[0]: raise ValueError('mismatched shape') if d is not None: assert d.ndim == 2 if not isinstance(d, cupyx.scipy.sparse.csr_matrix): raise TypeError('unsupported type (actual: {})'.format(type(d))) assert d.has_canonical_format if a.shape[0] != d.shape[0] or b.shape[1] != d.shape[1]: raise ValueError('mismatched shape') handle = device.get_cusparse_handle() m, k = a.shape _, n = b.shape if d is None: a, b = _cast_common_type(a, b) else: a, b, d = _cast_common_type(a, b, d) info = cusparse.createCsrgemm2Info() alpha = numpy.array(alpha, a.dtype).ctypes null_ptr = 0 if d is None: beta_data = null_ptr d_descr = MatDescriptor.create() d_nnz = 0 d_data = null_ptr d_indptr = null_ptr d_indices = null_ptr else: beta = numpy.array(beta, a.dtype).ctypes beta_data = beta.data d_descr = d._descr d_nnz = d.nnz d_data = d.data.data.ptr d_indptr = d.indptr.data.ptr d_indices = d.indices.data.ptr buff_size = _call_cusparse( 'csrgemm2_bufferSizeExt', a.dtype, handle, m, n, k, alpha.data, a._descr.descriptor, a.nnz, a.indptr.data.ptr, a.indices.data.ptr, b._descr.descriptor, b.nnz, b.indptr.data.ptr, b.indices.data.ptr, beta_data, d_descr.descriptor, d_nnz, d_indptr, d_indices, info) buff = cupy.empty(buff_size, numpy.int8) c_nnz = numpy.empty((), 'i') cusparse.setPointerMode(handle, cusparse.CUSPARSE_POINTER_MODE_HOST) c_descr = MatDescriptor.create() c_indptr = cupy.empty(m + 1, 'i') cusparse.xcsrgemm2Nnz( handle, m, n, k, a._descr.descriptor, a.nnz, a.indptr.data.ptr, a.indices.data.ptr, b._descr.descriptor, b.nnz, b.indptr.data.ptr, b.indices.data.ptr, d_descr.descriptor, d_nnz, d_indptr, d_indices, c_descr.descriptor, c_indptr.data.ptr, c_nnz.ctypes.data, info, buff.data.ptr) c_indices = cupy.empty(int(c_nnz), 'i') c_data = cupy.empty(int(c_nnz), a.dtype) _call_cusparse( 'csrgemm2', a.dtype, handle, m, n, k, alpha.data, a._descr.descriptor, a.nnz, a.data.data.ptr, a.indptr.data.ptr, a.indices.data.ptr, b._descr.descriptor, b.nnz, b.data.data.ptr, b.indptr.data.ptr, b.indices.data.ptr, beta_data, d_descr.descriptor, d_nnz, d_data, d_indptr, d_indices, c_descr.descriptor, c_data.data.ptr, c_indptr.data.ptr, c_indices.data.ptr, info, buff.data.ptr) c = cupyx.scipy.sparse.csr_matrix( (c_data, c_indices, c_indptr), shape=(m, n)) c._has_canonical_format = True cusparse.destroyCsrgemm2Info(info) return c
def csrgemm(a, b, transa=False, transb=False): """Matrix-matrix product for CSR-matrix. math:: C = op(A) op(B), Args: a (cupyx.scipy.sparse.csr_matrix): Sparse matrix A. b (cupyx.scipy.sparse.csr_matrix): Sparse matrix B. transa (bool): If ``True``, transpose of A is used. transb (bool): If ``True``, transpose of B is used. Returns: cupyx.scipy.sparse.csr_matrix: Calculated C. """ if not check_availability('csrgemm'): raise RuntimeError('csrgemm is not available.') assert a.ndim == b.ndim == 2 assert a.has_canonical_format assert b.has_canonical_format a_shape = a.shape if not transa else a.shape[::-1] b_shape = b.shape if not transb else b.shape[::-1] if a_shape[1] != b_shape[0]: raise ValueError('dimension mismatch') handle = device.get_cusparse_handle() m, k = a_shape n = b_shape[1] a, b = _cast_common_type(a, b) if a.nnz == 0 or b.nnz == 0: return cupyx.scipy.sparse.csr_matrix((m, n), dtype=a.dtype) op_a = _transpose_flag(transa) op_b = _transpose_flag(transb) nnz = numpy.empty((), 'i') cusparse.setPointerMode( handle, cusparse.CUSPARSE_POINTER_MODE_HOST) c_descr = MatDescriptor.create() c_indptr = cupy.empty(m + 1, 'i') cusparse.xcsrgemmNnz( handle, op_a, op_b, m, n, k, a._descr.descriptor, a.nnz, a.indptr.data.ptr, a.indices.data.ptr, b._descr.descriptor, b.nnz, b.indptr.data.ptr, b.indices.data.ptr, c_descr.descriptor, c_indptr.data.ptr, nnz.ctypes.data) c_indices = cupy.empty(int(nnz), 'i') c_data = cupy.empty(int(nnz), a.dtype) _call_cusparse( 'csrgemm', a.dtype, handle, op_a, op_b, m, n, k, a._descr.descriptor, a.nnz, a.data.data.ptr, a.indptr.data.ptr, a.indices.data.ptr, b._descr.descriptor, b.nnz, b.data.data.ptr, b.indptr.data.ptr, b.indices.data.ptr, c_descr.descriptor, c_data.data.ptr, c_indptr.data.ptr, c_indices.data.ptr) c = cupyx.scipy.sparse.csr_matrix( (c_data, c_indices, c_indptr), shape=(m, n)) c._has_canonical_format = True return c