def csr2coo(x, data, indices):
"""Converts a CSR-matrix to COO format.
Args:
x (cupyx.scipy.sparse.csr_matrix): A matrix to be converted.
data (cupy.ndarray): A data array for converted data.
indices (cupy.ndarray): An index array for converted data.
Returns:
cupyx.scipy.sparse.coo_matrix: A converted matrix.
"""
if not check_availability('csr2coo'):
raise RuntimeError('csr2coo is not available.')
handle = device.get_cusparse_handle()
m = x.shape[0]
nnz = len(x.data)
row = cupy.empty(nnz, 'i')
cusparse.xcsr2coo(
handle, x.indptr.data.ptr, nnz, m, row.data.ptr,
cusparse.CUSPARSE_INDEX_BASE_ZERO)
# data and indices did not need to be copied already
return cupyx.scipy.sparse.coo_matrix(
(data, (row, indices)), shape=x.shape)
def _csr_row_index(rows, Ap, Aj, Ax):
"""Populate indices and data arrays from the given row index
Args:
rows (cupy.ndarray): index array of rows to populate
Ap (cupy.ndarray): indptr array from input sparse matrix
Aj (cupy.ndarray): indices array from input sparse matrix
Ax (cupy.ndarray): data array from input sparse matrix
Returns:
Bp (cupy.ndarray): indptr array for output sparse matrix
Bj (cupy.ndarray): indices array of output sparse matrix
Bx (cupy.ndarray): data array of output sparse matrix
"""
row_nnz = cupy.diff(Ap)
Bp = cupy.empty(rows.size + 1, dtype=Ap.dtype)
Bp[0] = 0
cupy.cumsum(row_nnz[rows], out=Bp[1:])
nnz = int(Bp[-1])
out_rows = cupy.empty(nnz, dtype=numpy.int32)
# Build a COO row array from output CSR indptr.
# Calling backend cusparse API directly to avoid
# constructing a whole COO object.
handle = device.get_cusparse_handle()
cusparse.xcsr2coo(handle, Bp.data.ptr, nnz, Bp.size - 1, out_rows.data.ptr,
cusparse.CUSPARSE_INDEX_BASE_ZERO)
Bj, Bx = _csr_row_index_ker(out_rows, rows, Ap, Aj, Ax, Bp)
return Bp, Bj, Bx
def _csr_indptr_to_coo_rows(nnz, Bp):
out_rows = cupy.empty(nnz, dtype=numpy.int32)
# Build a COO row array from output CSR indptr.
# Calling backend cusparse API directly to avoid
# constructing a whole COO object.
handle = device.get_cusparse_handle()
cusparse.xcsr2coo(handle, Bp.data.ptr, nnz, Bp.size - 1, out_rows.data.ptr,
cusparse.CUSPARSE_INDEX_BASE_ZERO)
return out_rows
def _csr_indptr_to_coo_rows(nnz, Bp):
out_rows = cupy.empty(nnz, dtype=numpy.int32)
# Build a COO row array from output CSR indptr.
# Calling backend cusparse API directly to avoid
# constructing a whole COO object.
handle = device.get_cusparse_handle()
if runtime.is_hip and nnz == 0:
raise ValueError('hipSPARSE currently cannot handle '
'sparse matrices with null ptrs')
cusparse.xcsr2coo(handle, Bp.data.ptr, nnz, Bp.size - 1, out_rows.data.ptr,
cusparse.CUSPARSE_INDEX_BASE_ZERO)
return out_rows
def csc2coo(x, data, indices):
"""Converts a CSC-matrix to COO format.
Args:
x (cupyx.scipy.sparse.csc_matrix): A matrix to be converted.
data (cupy.ndarray): A data array for converted data.
indices (cupy.ndarray): An index array for converted data.
Returns:
cupyx.scipy.sparse.coo_matrix: A converted matrix.
"""
handle = device.get_cusparse_handle()
n = x.shape[1]
nnz = len(x.data)
col = cupy.empty(nnz, 'i')
cusparse.xcsr2coo(handle, x.indptr.data.ptr, nnz, n, col.data.ptr,
cusparse.CUSPARSE_INDEX_BASE_ZERO)
# data and indices did not need to be copied already
return cupyx.scipy.sparse.coo_matrix((data, (indices, col)), shape=x.shape)
def _csr_row_index(rows, Ap, Aj, Ax, Bp):
"""Populate indices and data arrays from the given row index
Args
rows : index array of rows to populate
Ap : indptr array from input sparse matrix
Aj : indices array from input sparse matrix
Ax : data array from input sparse matrix
Bp : indptr array for output sparse matrix
tpb : threads per block of row index kernel
Returns
Bj : indices array of output sparse matrix
Bx : data array of output sparse matrix
"""
nnz = int(Bp[-1])
Bj = cupy.empty(nnz, dtype=Aj.dtype)
Bx = cupy.empty(nnz, dtype=Ax.dtype)
out_rows = cupy.empty(nnz, dtype=rows.dtype)
# Build a COO row array from output CSR indptr.
# Calling backend cusparse API directly to avoid
# constructing a whole COO object.
handle = device.get_cusparse_handle()
cusparse.xcsr2coo(handle, Bp.data.ptr, nnz, Bp.size - 1, out_rows.data.ptr,
cusparse.CUSPARSE_INDEX_BASE_ZERO)
_csr_row_index_ker(out_rows,
rows,
Ap,
Aj,
Ax,
Bp,
Bj,
Bx,
size=out_rows.size)
return Bj, Bx