def cscsort(x): """Sorts indices of CSC-matrix in place. Args: x (cupyx.scipy.sparse.csc_matrix): A sparse matrix to sort. """ if not check_availability('cscsort'): raise RuntimeError('cscsort is not available.') nnz = x.nnz if nnz == 0: return handle = device.get_cusparse_handle() m, n = x.shape buffer_size = cusparse.xcscsort_bufferSizeExt( handle, m, n, nnz, x.indptr.data.ptr, x.indices.data.ptr) buf = cupy.empty(buffer_size, 'b') P = cupy.empty(nnz, 'i') data_orig = x.data.copy() cusparse.createIdentityPermutation(handle, nnz, P.data.ptr) cusparse.xcscsort( handle, m, n, nnz, x._descr.descriptor, x.indptr.data.ptr, x.indices.data.ptr, P.data.ptr, buf.data.ptr) _call_cusparse( 'gthr', x.dtype, handle, nnz, data_orig.data.ptr, x.data.data.ptr, P.data.ptr, cusparse.CUSPARSE_INDEX_BASE_ZERO)
def coosort(x, sort_by='r'): """Sorts indices of COO-matrix in place. Args: x (cupyx.scipy.sparse.coo_matrix): A sparse matrix to sort. sort_by (str): Sort the indices by row ('r', default) or column ('c'). """ if not check_availability('coosort'): raise RuntimeError('coosort is not available.') nnz = x.nnz if nnz == 0: return handle = device.get_cusparse_handle() m, n = x.shape buffer_size = cusparse.xcoosort_bufferSizeExt( handle, m, n, nnz, x.row.data.ptr, x.col.data.ptr) buf = cupy.empty(buffer_size, 'b') P = cupy.empty(nnz, 'i') data_orig = x.data.copy() cusparse.createIdentityPermutation(handle, nnz, P.data.ptr) if sort_by == 'r': cusparse.xcoosortByRow( handle, m, n, nnz, x.row.data.ptr, x.col.data.ptr, P.data.ptr, buf.data.ptr) elif sort_by == 'c': cusparse.xcoosortByColumn( handle, m, n, nnz, x.row.data.ptr, x.col.data.ptr, P.data.ptr, buf.data.ptr) else: raise ValueError("sort_by must be either 'r' or 'c'") _call_cusparse( 'gthr', x.dtype, handle, nnz, data_orig.data.ptr, x.data.data.ptr, P.data.ptr, cusparse.CUSPARSE_INDEX_BASE_ZERO) if sort_by == 'c': # coo is sorted by row first x._has_canonical_format = False