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