IntelPython / dpnp / 42105526684

84%

push

github

Sort `csr_matrix` indices once on a combined row and column key (#3134)

`dpnp.scipy.sparse.csr_matrix.sort_indices` sorted the stored entries by
(row, column) with two stable argsorts: one on the column indices, then
one on the permuted row ids. This PR sorts once on the int64 key `row *
ncols + col`. CSR rows are already in order, so the key is nearly
sorted, and one merge sort of it is much cheaper than two sorts of fully
permuted data.

* The per-row offsets come from `repeat(arange(0, nrows * ncols, ncols),
row_lengths)`, so only one elementwise add runs before the sort.
Computing `row_ids * ncols + indices` instead adds kernels that made
small matrices slower than before on GPU (0.71-0.95x of master below
about 3·10⁴ stored values).
* The offsets are int64 even for int32 indices, so the key can't
overflow. `has_sorted_indices` runs `check_format` first, which bounds
the column indices to `[0, ncols)`.
* When `nrows * ncols >= 2**63` the key doesn't fit int64, and the
previous two-pass sort is kept. Its row ids are now int64 too.
* `kind="stable"` (merge sort for int64) stays: radix sort on the int64
key was slower on both devices measured.

Two new tests: an int32-index matrix whose key overflows int32 (shape
3001 x 10⁶), and shapes just below, at and above the int64 key limit,
the last of which takes the fallback. They fail when the offsets use the
index dtype or when the fallback is removed. Both allocate only a few
elements whatever the shape, so they aren't marked `slow`.

### Performance

`sort_indices` alone (construction excluded), master vs this PR loaded
in the same process, median time, on a Panther Lake H machine: the Xe3
iGPU on Level Zero and the Core Ultra 7 366H on OpenCL. Each matrix has
8 or about 1024 stored values per row, with columns shuffled within each
row; results are identical to master in every case.

| nnz | shape | index | GPU master, ms | GPU PR, ms | GPU speedup | CPU
master, ms | CPU PR, ms | CPU speedup |
| ---: |... (continued)

1792 of 2998 branches covered (59.77%)

Branch coverage included in aggregate %.

8 of 8 new or added lines in 1 file covered. (100.0%)

29619 of 34220 relevant lines covered (86.55%)

39123.81 hits per line