Sort csr_matrix indices once on a combined row and column key

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)

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]

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())