Repository navigation
Add cho_factor and cho_solve to dpnp.scipy.linalg - #3116
antonwolfy wants to merge 17 commits into
Conversation
|
View rendered docs @ https://intelpython.github.io/dpnp/pull/3116/index.html |
|
Array API standard conformance tests for dpnp=0.22.0dev0=np2py314h8d9cdd5_20 ran successfully. |
Implement SciPy-compatible `dpnp.scipy.linalg.cho_factor` and `dpnp.scipy.linalg.cho_solve` with batched input support (gh-3096). - `cho_factor` uses the existing `potrf` / `potrf_batch` bindings. Only the referenced triangle is copied into a zeroed C-order buffer, which matches the SciPy 1.18 output in a single pass. The input is factorized in place for `overwrite_a=True` when it is 2-D, writable, of the result dtype, and C- or F-contiguous. - `cho_solve` uses a new LAPACK `potrs` binding for 2-D input and a new BLAS `trsm_batch` binding (two triangular solves) for batched input, since oneMKL `potrs_batch` is much slower. Batch dimensions are broadcast as in `lu_solve`. - Add tests and API reference entries.
A C-order factor is passed to oneMKL transposed, so the solves are with conj(A), and `b` and the solution are conjugated to compensate. This is only needed for a complex factor. Decide it from the dtype of `c` rather than the result dtype, so a real factor with a complex right-hand side no longer conjugates `b` and the solution. Add a test for a real factor with complex right-hand sides.
Add `tests/cupyx_tests/linalg_tests/test_solve.py` from CuPy as `dpnp/tests/third_party/cupyx/linalg_tests/test_solve.py`, with `TestChosolve` adapted to `dpnp.scipy.linalg.cho_solve`: - pass `check_finite=False`, since dpnp checks the whole factor for infs and NaNs, as SciPy does, and the tests put NaNs in the unused triangle; - pass a batch of vectors with a trailing axis, since dpnp follows SciPy's `(..., n, k)` rule for a batched right-hand side. The `invh` tests are skipped, as `invh` is not supported.
- Add See Also link from dpnp.linalg.cholesky to dpnp.scipy.linalg.cho_factor - Fix the upper-triangular example output in dpnp.linalg.cholesky - Remove stale linalg.lu_factor and linalg.lu_solve entries from linalg.rst - Look up the order manager in _copy_rhs and _cho_solve_trsm instead of passing it
- Do not broadcast a 2-D factor `c` over a batched right-hand side: fold the batch into the RHS columns for one `potrs` call, or, for a real type on GPU with a small batch (batch * nrhs > 16, batch <= 32), where oneMKL potrs slows down above 16 columns, share the factor by all batch elements in `trsm_batch` - Let `_trsm_batch` take a 2-D triangular matrix shared by the batch (stride 0) - Submit both `check_finite` reductions before waiting for the first result - Accept the batch-shaped `lower` array returned by SciPy >= 1.18 in `cho_solve`, if all its values are equal
- `bool()` of a non-empty list or tuple is always True, so `cho_solve((c, [False]), b)` used the wrong triangle: convert a sequence like an array, as SciPy's `cho_solve` does, and do the same in `cho_factor` - Use finite values in `TestChoSolve.test_empty`: `check_finite` reads `c` or `b` when it is not empty, so uninitialized data could fail it - Add the changelog entry
410539a to
6b3035e
Compare
…ciPy differences - `cho_factor` with `overwrite_a=True` and `cho_solve` with `overwrite_b=True` returned a `usm_ndarray` input as is: wrap it with `dpnp.get_result_array` - Reject a 0-d right-hand side in the `_potrs` binding instead of reading its shape - Document that `cho_factor` returns `lower` as a bool and that `cho_solve` supports a zero-size batch, unlike SciPy
- Move the dispatch rule for a 2-D factor with a batched right-hand side to `_fold_rhs_batch`, so tests can force each path on any device - Sample integers from [-2, 2) in `_make_hpd_np`: [-1, 1) truncated them to zeros, so integer test matrices were diagonal - Test a batched column-major factor, an empty `lower` array, the dispatch boundaries, a non-finite value in the triangle `cho_factor` does not reference, an in-place 1-D `b`, a 1-D `c` and mixed queues, and match the error messages - Drop the direct test of the `_potrs` binding
- List the `ValueError` cases of `cho_factor` next to `LinAlgError` - Add a Raises section to `cho_solve` - Mention the sync to read a dpnp array `lower` in `cho_solve`
- Catch `oneapi::mkl::exception` after the LAPACK one, so an error not from LAPACK (e.g. `unimplemented` or a failed allocation) also frees the scratchpad - Include the standard headers used and drop the unused `linalg_exceptions.hpp`
A `c` of shape (1, n, n) or (1, 1, n, n) with a batched right-hand side was broadcast to the batch and copied per batch element; solve with it as with a 2-D factor instead
- View a 0-D or 1-D `a` as a matrix in `cho_factor`, as SciPy does, and report its shape as `a1.shape` - Raise `ValueError` "The factored matrix c is not square." in `cho_solve` for `c` with fewer than two dimensions - Report the shapes of a single system for incompatible `c` and `b`, also for batched input
- Copy `b` into the solve buffer through a view in its batch shape, so a broadcast `b` is not first reshaped into a full copy - Restore `_align_lu_solve_broadcast` as on master: `cho_solve` checks the shapes before calling it - Include the standard headers used in `potrs.hpp`, `trsm.hpp` and `trsm_batch.cpp`
There was a problem hiding this comment.
🟡 Changes recommended
Half-precision inputs can reach unsupported Cholesky dispatches, and overwriting a mixed-USM right-hand side can return a solution with the wrong USM type.
2 open issues
What changed in this PR
Adds SciPy-compatible Cholesky factorization and solving to dpnp.scipy.linalg, enabling positive-definite systems to use Cholesky rather than general LU solves.
Changes:
- Adds
cho_factorandcho_solvewith batching, broadcasting, overwrite options, and finite-value checks. - Introduces oneMKL
potrsand batched triangular-solve bindings. - Adds numerical, queue, and USM tests, plus API documentation and a changelog entry.
| File | Description |
|---|---|
dpnp/tests/third_party/cupyx/linalg_tests/test_solve.py |
Adapts CuPy Cholesky-solve tests; skips unsupported invh. |
dpnp/tests/third_party/cupyx/linalg_tests/__init__.py |
Adds the test package marker. |
dpnp/tests/test_usm_type.py |
Checks result USM allocation types. |
dpnp/tests/test_sycl_queue.py |
Checks queue preservation and mixed-queue rejection. |
dpnp/tests/test_linalg.py |
Tests numerical results, batching, layouts, and validation. |
dpnp/scipy/linalg/_utils.py |
Implements factorization and optimized solve paths. |
dpnp/scipy/linalg/_decomp_cholesky.py |
Defines and documents the public Cholesky APIs. |
dpnp/scipy/linalg/__init__.py |
Exports both new functions. |
dpnp/linalg/dpnp_iface_linalg.py |
Adds a documentation cross-reference and corrects an example. |
dpnp/backend/extensions/lapack/types_matrix.hpp |
Defines supported potrs types. |
dpnp/backend/extensions/lapack/potrs.hpp |
Declares the potrs binding. |
dpnp/backend/extensions/lapack/potrs.cpp |
Implements validated oneMKL Cholesky solves. |
dpnp/backend/extensions/lapack/lapack_py.cpp |
Exposes potrs to Python. |
dpnp/backend/extensions/lapack/CMakeLists.txt |
Includes potrs in the build. |
dpnp/backend/extensions/blas/types_matrix.hpp |
Defines supported batched triangular-solve types. |
dpnp/backend/extensions/blas/trsm.hpp |
Declares the batched triangular-solve binding. |
dpnp/backend/extensions/blas/trsm_batch.cpp |
Implements batched solves with shared-factor support. |
dpnp/backend/extensions/blas/CMakeLists.txt |
Includes batched triangular solves in the build. |
dpnp/backend/extensions/blas/blas_py.cpp |
Exposes trsm_batch to Python. |
doc/reference/scipy_linalg.rst |
Lists the new Cholesky functions. |
doc/reference/linalg.rst |
Removes stale LU API entries. |
CHANGELOG.md |
Records the new public APIs. |
Review effort: Balanced
With `overwrite_b=True`, a right-hand side whose USM type differs from the coerced one of `c` and `b` was solved in place, so the solution kept the USM type of `b`. Copy it in that case, as without `overwrite_b`.
|
| Shape | dtype | NumPy CPU | dpnp SVD (today) | A: host solve | B: device LU | C: device Cholesky | C vs B | C vs NumPy |
|---|---|---|---|---|---|---|---|---|
| 500,000×100 | f32 | 0.072 | 0.784 | 0.080 | 0.093 | 0.096 | 0.96x | 0.75x |
| 50,000×1000 | f32 | 0.199 | 1.079 | 0.105 | 0.080 | 0.076 | 1.06x | 2.6x |
| 50,000×1000 | f64 | 0.251 | 1.687 | 0.249 | 0.256 | 0.237 | 1.08x | 1.06x |
| 463,810×990 | f32 | 1.194 | - | 0.771 | 0.758 | 0.765 | 0.99x | 1.56x |
| 2,500×20,000 | f32 | 0.278 | 5.696 | 0.176 | 0.104 | 0.095 | 1.09x | 2.9x |
| 2,500×20,000 | f64 | 0.723 | 6.595 | 0.498 | 0.464 | 0.444 | 1.05x | 1.63x |
| 5,000×40,000 | f32 | 1.969 | - | 0.766 | 0.554 | 0.538 | 1.03x | 3.7x |
| 10,000×40,000 | f32 | 7.210 | - | 2.386 | 1.834 | 1.473 | 1.25x | 4.9x |
| 100,000×5,000 | f32 | 4.360 | - | 1.451 | 1.358 | 1.280 | 1.06x | 3.4x |
The last two shapes are not in that issue; they are added to show larger systems.
Solve step alone (the n×n system inside fit), ms:
| n (shape) | dtype | B: LU | C: Cholesky | C vs B |
|---|---|---|---|---|
| 100 (500,000×100) | f32 | 0.64 | 3.2 | 0.2x |
| 1000 (50,000×1000) | f32 | 21.0 | 11.6 | 1.8x |
| 1000 (50,000×1000) | f64 | 30.2 | 14.0 | 2.2x |
| 2500 (2,500×20,000) | f32 | 41.0 | 19.6 | 2.1x |
| 2500 (2,500×20,000) | f64 | 175.8 | 82.2 | 2.1x |
| 5000 (5,000×40,000) | f32 | 165.2 | 95.6 | 1.7x |
| 5000 (100,000×5,000) | f32 | 189.2 | 102.1 | 1.9x |
| 10000 (10,000×40,000) | f32 | 588.4 | 282.4 | 2.1x |
Takeaways:
- Replacing SVD is the large gain (7–60x), and all three variants get it.
- The Cholesky solve is about 2x faster than LU from n=1000 up. In
fitthe change is between −4% and +9% on the shapes from the issue, and +25% on 10,000×40,000, since the Gram matmul dominates; the gain grows with the size of the system. - Cholesky on the device is 1.4–1.9x faster than the host-solve variant A on the f32 dual shapes and on 50,000×1000 f32, and 1.05–1.1x in f64.
- For n=100, Cholesky is slower than LU because of the small-n
potrflatency on client GPUs (MKLD-20346). - For tall data, most of the fit time is scikit-learn's own validation and centering on dpnp, about 0.3 ms per MiB of
X(517 ms of the 765 ms fit at 463,810×990). That is worth a separate look. - The remaining part of Add Cholesky-based linear solvers to dpnp.scipy.linalg: cho_factor, cho_solve, solve_triangular, solve #3096,
solve(assume_a="pos"), is what a SciPy-style caller would use to reach this Cholesky path.

Add SciPy-compatible
dpnp.scipy.linalg.cho_factoranddpnp.scipy.linalg.cho_solve, including batched input as in SciPy 1.18 (#3096).Implementation
cho_factoruses the existingpotrfandpotrf_batchbindings. Only the referenced triangle is copied, in one pass, into a zeroed C-order buffer, which matches the SciPy 1.18 output. Withoverwrite_a=True, a 2-D, writable input of the result dtype that is C- or F-contiguous is factored in place. As in SciPy, a 0-D or 1-D input is viewed as a matrix.cho_solveuses a new LAPACKpotrsbinding for a 2-D factor. For a batch of factors it calls a new BLAStrsm_batchbinding twice, since oneMKLpotrs_batchis much slower (MKLD-20341).potrscall. For a real type on GPU with more than 16 folded columns and a batch of at most 32, the factor is shared by the batch intrsm_batch(stride 0) instead, because double-precisiontrsmon Xe GPUs slows down by up to 3.6x above 16 columns.lu_solve, and a broadcast right-hand side is copied without being materialized first.lowercan also be the batch-shaped array SciPy 1.18 returns, or a sequence, if all its values are equal.check_finite=True, both finiteness checks ofcho_solveare submitted before the first host sync.Differences from SciPy
Documented in the docstrings and covered by tests:
cho_factorreturnsloweras a Python bool; SciPy 1.18 returns an array tiled to the batch shape.cho_solvesupports a zero-size batch; SciPy raises.As in the existing
lu_factoranddpnp.linalg.cholesky, integer and boolean inputs are cast to the default floating type of the device; SciPy 1.18 uses float32 for boolean and up to 16-bit integer inputs.Other changes
dpnp.linalg.cholesky: a See Also link todpnp.scipy.linalg.cho_factorand a fixed example output forupper=True.doc/reference/linalg.rst: removed the stalelinalg.lu_factorandlinalg.lu_solveentries.tests/cupyx_tests/linalg_tests/test_solve.pyfrom CuPy, withTestChosolveadapted todpnp.scipy.linalg.cho_solveand theinvhtests skipped.Testing
test_linalg.py,test_sycl_queue.pyandtest_usm_type.py, plus the CuPy tests.lower, batching and broadcasting,overwrite_*andcheck_finite) finds no differences. Earlier revisions were also tested on PVC, Lunar Lake and Arc B580.Performance
cho_factor+cho_solveagainstdpnp.linalg.solve, median times, measured on an earlier revision of this branch:cho_solvevslu_solveOn client and Battlemage GPUs, the Cholesky path is slower than LU for small single matrices (n=256: 0.5–0.6x) and for batches of larger matrices such as (100, 128, 128) (0.24–0.5x), because of oneMKL issues in
potrfandpotrf_batch.