diff --git a/CHANGELOG.md b/CHANGELOG.md index e1ddf041686..42347a46483 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +* Improved performance of `dpnp.scipy.sparse.csr_matrix.sort_indices` by sorting once on a combined row and column key instead of in two passes [#3134](https://github.com/IntelPython/dpnp/pull/3134) + ### Deprecated * Deprecated `dpnp.dot` and `dpnp.ndarray.dot` when `a.ndim >= 2` and `b.ndim > 2`, which place the leading dimensions of `b` in the middle of the result; use `dpnp.tensordot(a, b, axes=(-1, -2))` or `dpnp.matmul` instead [#3124](https://github.com/IntelPython/dpnp/pull/3124) diff --git a/dpnp/scipy/sparse/_csr.py b/dpnp/scipy/sparse/_csr.py index ee329623185..73d1970d097 100644 --- a/dpnp/scipy/sparse/_csr.py +++ b/dpnp/scipy/sparse/_csr.py @@ -385,20 +385,25 @@ def sort_indices(self): self._has_sorted_indices = True return - nrows = self._shape[0] + nrows, ncols = self._shape row_lengths = self._indptr[1:] - self._indptr[:-1] - row_ids = _dpnp.repeat( - _dpnp.arange( - nrows, - dtype=indices.dtype, - usm_type=indices.usm_type, - sycl_queue=indices.sycl_queue, - ), - row_lengths, - ) - # Lexsort by (row, col) via two stable passes. - order = _dpnp.argsort(indices, kind="stable") - order = order[_dpnp.argsort(row_ids[order], kind="stable")] + # int64 even for int32 indices, so neither branch can overflow. + common = { + "dtype": _dpnp.int64, + "usm_type": indices.usm_type, + "sycl_queue": indices.sycl_queue, + } + if nrows * ncols < 2**63: + # One sort on row * ncols + col; check_format bounded the indices. + row_keys = _dpnp.repeat( + _dpnp.arange(0, nrows * ncols, ncols, **common), row_lengths + ) + order = _dpnp.argsort(row_keys + indices, kind="stable") + else: + # The key would overflow int64: lexsort in two stable passes. + row_ids = _dpnp.repeat(_dpnp.arange(nrows, **common), row_lengths) + order = _dpnp.argsort(indices, kind="stable") + order = order[_dpnp.argsort(row_ids[order], kind="stable")] self._data = self._data[order] self._indices = self._indices[order] diff --git a/dpnp/tests/test_scipy_sparse_linalg.py b/dpnp/tests/test_scipy_sparse_linalg.py index cd994ffe985..d2af46262be 100644 --- a/dpnp/tests/test_scipy_sparse_linalg.py +++ b/dpnp/tests/test_scipy_sparse_linalg.py @@ -1587,6 +1587,38 @@ def test_sort_indices_idempotent(self): assert list(dpnp.asnumpy(m.indices)) == idx_after assert list(dpnp.asnumpy(m.data)) == data_after + def test_sort_indices_int32_key_overflow(self): + # Indices fit int32, but row * ncols + col does not. + data = dpnp.array([1.0, 2.0, 3.0, 4.0], dtype=dpnp.float32) + indices = dpnp.array([7, 5, 9, 2], dtype=dpnp.int32) + indptr = numpy.full(3002, 2, dtype=numpy.int32) + indptr[:2] = 0 + indptr[-1] = 4 + m = csr_matrix( + (data, indices, dpnp.asarray(indptr)), shape=(3001, 10**6) + ) + m.sort_indices() + assert m.indices.dtype == dpnp.int32 + assert list(dpnp.asnumpy(m.indices)) == [5, 7, 2, 9] + assert list(dpnp.asnumpy(m.data)) == [2.0, 1.0, 4.0, 3.0] + + @pytest.mark.parametrize( + "shape", + [(7, (2**63 - 1) // 7), (2, 2**62), (3, 2**62)], + ids=["below-int64-key", "at-int64-key", "above-int64-key"], + ) + def test_sort_indices_huge_shape(self, shape): + nrows, ncols = shape + data = dpnp.array([1.0, 2.0, 3.0, 4.0], dtype=dpnp.float32) + indices = dpnp.array([ncols - 1, 5, ncols - 2, 0], dtype=dpnp.int64) + indptr = numpy.full(nrows + 1, 2, dtype=numpy.int64) + indptr[0] = 0 + indptr[-1] = 4 + m = csr_matrix((data, indices, dpnp.asarray(indptr)), shape=shape) + m.sort_indices() + assert list(dpnp.asnumpy(m.indices)) == [5, ncols - 1, 0, ncols - 2] + assert list(dpnp.asnumpy(m.data)) == [2.0, 1.0, 4.0, 3.0] + def test_check_sorted_detects_sorted_components(self): # Already-sorted components; lazy _check_sorted must report True. data = dpnp.array([1.0, 2.0, 3.0, 4.0], dtype=dpnp.default_float_type())