Repository navigation
Add cho_factor and cho_solve to dpnp.scipy.linalg - #3116
antonwolfy wants to merge 25 commits into
Conversation
|
View rendered docs @ https://intelpython.github.io/dpnp/pull/3116/index.html |
|
Array API standard conformance tests for dpnp=0.22.0dev1=np2py314h8d9cdd5_37 ran successfully. |
410539a to
6b3035e
Compare
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
|
| 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.
There was a problem hiding this comment.
🔵 Needs a closer look
The new asynchronous potrs/trsm_batch bindings need device-backed oneMath/CUDA/HIP validation beyond the supplied Intel results, which this source-only environment cannot provide.
No open issues 🎉
2 resolved since last review
Review effort: Balanced
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
…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`
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`.
…patch - query the potrs scratchpad size inside the try and raise if its allocation fails - drop the unused unit_diag option of _trsm_batch - use the shared-factor trsm_batch path only on Intel GPUs - note the measured potrs_batch slowdown next to the internal ticket
38ec63c to
20b0f45
Compare
oneMKL reports a zero potrs scratchpad, so the clean-up only freed a nullptr. keep_args_alive now depends on the potrs event directly.
| # bool() of a non-empty sequence is always True | ||
| lower = numpy.asarray(lower) | ||
| if isinstance(lower, numpy.ndarray) and lower.size != 1: | ||
| if lower.size == 0: |
There was a problem hiding this comment.
An empty lower is read as False here, so cho_solve((c, []), b) and cho_solve((c, dpnp.array([])), b) silently solve with the upper triangle. SciPy 1.18.1 raises for both:
IndexError: _flapack.dpotrs() 1st keyword (lower) can't be converted to int
Non-empty arrays of equal values (e.g. [True, True]) are accepted by both, so only the empty case differs. Could this raise instead, with a test? The lower bullet in the PR description should then be updated to match.
There was a problem hiding this comment.
Confirmed, an empty lower is now rejected with a ValueError in cho_factor and cho_solve. The one exception is cho_solve with a zero-size batch of factors, where SciPy 1.18 tiles lower to the empty batch shape, so an empty lower is valid there (test_lower_array_empty). Tests added for both functions; the lower bullet in the PR description will be updated.
There was a problem hiding this comment.
A string has the same problem: "False" is read as True, so the solve uses the wrong triangle, while SciPy raises ValueError:
import dpnp
from dpnp.scipy.linalg import cho_factor, cho_solve
a = dpnp.array([[4.0, 2.0], [2.0, 3.0]])
b = dpnp.array([1.0, 2.0])
c, low = cho_factor(a) # upper factor, the lower triangle is zero
print(cho_solve((c, low), b), cho_solve((c, "False"), b))
# GPU and CPU: [-0.125 0.75 ] [0.25 1. ]; SciPy raises ValueError for "False"When the empty case raises, test_lower_array_empty (an empty lower with a zero-size batch) should still pass.
There was a problem hiding this comment.
Fixed: lower is now converted with dpnp.asnumpy and anything that is not a bool, integer or float array (strings, None, complex, object) raises ValueError, so "False" no longer selects the wrong triangle. The empty case with a zero-size batch still passes (test_lower_array_empty). Added test_lower_invalid_type to TestChoFactor and TestChoSolve.
|
|
||
| linalg.cholesky | ||
| linalg.qr | ||
| linalg.lu_factor |
There was a problem hiding this comment.
Removing the stale linalg.lu_factor/linalg.lu_solve entries (here and below) and fixing the upper=True example output of dpnp.linalg.cholesky both look right, but they are unrelated to cho_factor/cho_solve. Per AGENTS.md ("One PR, one change"), could they move to a separate PR? The See Also link to cho_factor belongs here. The "Other changes" section of the PR description would need to be updated accordingly.
There was a problem hiding this comment.
It's a minor doc improvement and I don't think we need a separate PR on that.
The PR description updated.
| if (is_exception_caught) // an unexpected error occurs | ||
| { | ||
| if (scratchpad != nullptr) { | ||
| dpnp::tensor::alloc_utils::sycl_free_noexcept(scratchpad, exec_q); |
There was a problem hiding this comment.
If mkl_lapack::potrs throws after it has already enqueued work (e.g. a sycl::exception from a later internal submission), the scratchpad is freed here while a kernel may still be using it. Usually the call throws before submitting anything, so this is low probability, and it's the same pattern as in the other LAPACK extensions (getrs.cpp, potrf_batch.cpp, ...). Still, would waiting on the queue before the free on this path be safer, possibly as a follow-up across all of them?
There was a problem hiding this comment.
That sound valid and it would be safer to call wait and no impact on legacy performance.
I will open the follow-up PR on that.
* Raise `ValueError` for an empty `lower` unless `cho_solve` gets a zero-size batch of factors * Copy `b` in `cho_solve` when `overwrite_b=True` and `b` overlaps the factor * Return early from `trsm_batch` for zero `m`, `n` or batch size * Document the array-like `lower` of `cho_factor` and test its errors
…st the original factor' Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Strings, `None` and complex values were read as `True` or silently truncated, SciPy raises for them. Also add a test for a strided factor overlapping `b` with `overwrite_b=True`.


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. An emptylowerraisesValueError, as in SciPy, except incho_solvefor a zero-size batch of factors.overwrite_b=True,cho_solvesolves in a copy whenboverlaps the factor (SciPy overwrites the factor).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.lowerwith unequal values raisesValueErrorincho_factorandcho_solve; SciPy 1.18 accepts it.cho_solvewithoverwrite_b=Trueandboverlapping the factor leaves the factor unchanged; SciPy overwrites it.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.