diff --git a/doc/reference/linalg.rst b/doc/reference/linalg.rst index 1072b5a019fe..9f293f3abd8c 100644 --- a/doc/reference/linalg.rst +++ b/doc/reference/linalg.rst @@ -43,7 +43,6 @@ Decompositions linalg.cholesky linalg.qr - linalg.lu_factor linalg.svd linalg.svdvals @@ -86,7 +85,6 @@ Solving linear equations linalg.solve linalg.tensorsolve linalg.lstsq - linalg.lu_solve linalg.inv linalg.pinv linalg.tensorinv diff --git a/doc/reference/scipy_linalg.rst b/doc/reference/scipy_linalg.rst index c45fa55de721..3a1b8e27514c 100644 --- a/doc/reference/scipy_linalg.rst +++ b/doc/reference/scipy_linalg.rst @@ -12,6 +12,8 @@ Decompositions :toctree: generated/ :nosignatures: + cho_factor + cho_solve lu lu_factor lu_solve diff --git a/dpnp/backend/extensions/blas/CMakeLists.txt b/dpnp/backend/extensions/blas/CMakeLists.txt index 872dea406ed3..44bb8be47083 100644 --- a/dpnp/backend/extensions/blas/CMakeLists.txt +++ b/dpnp/backend/extensions/blas/CMakeLists.txt @@ -34,6 +34,7 @@ set(_module_src ${CMAKE_CURRENT_SOURCE_DIR}/gemm_batch.cpp ${CMAKE_CURRENT_SOURCE_DIR}/gemv.cpp ${CMAKE_CURRENT_SOURCE_DIR}/syrk.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/trsm_batch.cpp ) pybind11_add_module(${python_module_name} MODULE ${_module_src}) diff --git a/dpnp/backend/extensions/blas/blas_py.cpp b/dpnp/backend/extensions/blas/blas_py.cpp index 28714194abb2..b976f2d02a9b 100644 --- a/dpnp/backend/extensions/blas/blas_py.cpp +++ b/dpnp/backend/extensions/blas/blas_py.cpp @@ -43,6 +43,7 @@ #include "gemm.hpp" #include "gemv.hpp" #include "syrk.hpp" +#include "trsm.hpp" namespace blas_ns = dpnp::extensions::blas; namespace py = pybind11; @@ -58,6 +59,7 @@ void init_dispatch_vectors_tables(void) blas_ns::init_gemm_dispatch_table(); blas_ns::init_gemv_dispatch_vector(); blas_ns::init_syrk_dispatch_vector(); + blas_ns::init_trsm_batch_dispatch_vector(); } static dot_impl_fn_ptr_t dot_dispatch_vector[dpnp_td_ns::num_types]; @@ -162,6 +164,17 @@ PYBIND11_MODULE(_blas_impl, m, py::mod_gil_not_used()) py::arg("depends") = py::list()); } + { + m.def("_trsm_batch", &blas_ns::trsm_batch, + "Call `trsm_batch` from oneMKL BLAS library to solve " + "op(A) X = B for a batch of triangular matrices A " + "(a 2-D A is shared by the batch). " + "trans_op picks op(A): 0=N, 1=T, 2=C.", + py::arg("sycl_queue"), py::arg("matrixA"), py::arg("matrixB"), + py::arg("upper_lower"), py::arg("trans_op") = 0, + py::arg("unit_diag") = false, py::arg("depends") = py::list()); + } + { m.def( "_using_onemath", diff --git a/dpnp/backend/extensions/blas/trsm.hpp b/dpnp/backend/extensions/blas/trsm.hpp new file mode 100644 index 000000000000..5f4682eb4087 --- /dev/null +++ b/dpnp/backend/extensions/blas/trsm.hpp @@ -0,0 +1,48 @@ +//***************************************************************************** +// Copyright (c) 2026, Intel Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// - Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// - Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +// THE POSSIBILITY OF SUCH DAMAGE. +//***************************************************************************** + +#pragma once + +#include +#include + +#include "dpnp4pybind11.hpp" + +namespace dpnp::extensions::blas +{ +extern std::pair + trsm_batch(sycl::queue &exec_q, + const dpnp::tensor::usm_ndarray &matrixA, + const dpnp::tensor::usm_ndarray &matrixB, + const std::int8_t upper_lower, + const int trans_op, + const bool unit_diag, + const std::vector &depends); + +extern void init_trsm_batch_dispatch_vector(void); +} // namespace dpnp::extensions::blas diff --git a/dpnp/backend/extensions/blas/trsm_batch.cpp b/dpnp/backend/extensions/blas/trsm_batch.cpp new file mode 100644 index 000000000000..b0cf8328dfda --- /dev/null +++ b/dpnp/backend/extensions/blas/trsm_batch.cpp @@ -0,0 +1,265 @@ +//***************************************************************************** +// Copyright (c) 2026, Intel Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// - Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// - Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +// THE POSSIBILITY OF SUCH DAMAGE. +//***************************************************************************** + +#include + +#include + +// utils extension header +#include "ext/common.hpp" + +// dpnp tensor headers +#include "utils/memory_overlap.hpp" +#include "utils/output_validation.hpp" +#include "utils/type_utils.hpp" + +#include "trsm.hpp" +#include "types_matrix.hpp" + +namespace dpnp::extensions::blas +{ +namespace mkl_blas = oneapi::mkl::blas; +namespace py = pybind11; +namespace type_utils = dpnp::tensor::type_utils; + +using ext::common::init_dispatch_vector; + +typedef sycl::event (*trsm_batch_impl_fn_ptr_t)( + sycl::queue &, + const oneapi::mkl::uplo, + const oneapi::mkl::transpose, + const oneapi::mkl::diag, + const std::int64_t, // m + const std::int64_t, // n + const char *, // a + const std::int64_t, // lda + const std::int64_t, // stride_a + char *, // b + const std::int64_t, // ldb + const std::int64_t, // stride_b + const std::int64_t, // batch_size + const std::vector &); + +static trsm_batch_impl_fn_ptr_t + trsm_batch_dispatch_vector[dpnp_td_ns::num_types]; + +template +static sycl::event trsm_batch_impl(sycl::queue &exec_q, + const oneapi::mkl::uplo upper_lower, + const oneapi::mkl::transpose trans, + const oneapi::mkl::diag unit_diag, + const std::int64_t m, + const std::int64_t n, + const char *matrixA, + const std::int64_t lda, + const std::int64_t stride_a, + char *matrixB, + const std::int64_t ldb, + const std::int64_t stride_b, + const std::int64_t batch_size, + const std::vector &depends) +{ + type_utils::validate_type_for_device(exec_q); + + const T *a = reinterpret_cast(matrixA); + T *b = reinterpret_cast(matrixB); + + std::stringstream error_msg; + bool is_exception_caught = false; + + sycl::event trsm_batch_event; + try { + // Release GIL to avoid serialization of host task submissions + // to the same queue in OneMKL + py::gil_scoped_release lock{}; + + trsm_batch_event = mkl_blas::column_major::trsm_batch( + exec_q, + oneapi::mkl::side::left, // op(A) is on the left: op(A) X = B + upper_lower, // Whether A is upper or lower triangular. + trans, // Defines op(A): 'N' for A, 'T' for A^T, or 'C' + // for A^H. + unit_diag, // Whether A is assumed to be unit triangular. + m, // Number of rows of B, and the order of A. + n, // Number of columns of B. + T(1), // Scaling factor for B. + a, // Pointer to the batch of triangular matrices A. + lda, // Leading dimension of each matrix A. + stride_a, // Stride between consecutive A matrices. + b, // Pointer to the batch of matrices B, overwritten + // by the solutions X. + ldb, // Leading dimension of each matrix B. + stride_b, // Stride between consecutive B matrices. + batch_size, // Total number of matrices in the batch. + depends); + } catch (oneapi::mkl::exception const &e) { + error_msg << "Unexpected MKL exception caught during trsm_batch() " + "call:\nreason: " + << e.what(); + is_exception_caught = true; + } catch (sycl::exception const &e) { + error_msg + << "Unexpected SYCL exception caught during trsm_batch() call:\n" + << e.what(); + is_exception_caught = true; + } + + if (is_exception_caught) // an unexpected error occurs + { + throw std::runtime_error(error_msg.str()); + } + + return trsm_batch_event; +} + +std::pair + trsm_batch(sycl::queue &exec_q, + const dpnp::tensor::usm_ndarray &matrixA, + const dpnp::tensor::usm_ndarray &matrixB, + const std::int8_t upper_lower, + const int trans_op, + const bool unit_diag, + const std::vector &depends) +{ + const int a_nd = matrixA.get_ndim(); + const int b_nd = matrixB.get_ndim(); + + // batch is the last dimension: (m, m, batch) and (m, n, batch), + // a 2-D (m, m) matrix A is shared by all right-hand sides + if ((a_nd != 2 && a_nd != 3) || b_nd != 3) { + throw py::value_error( + "The triangular and right-hand sides arrays must be " + "2- or 3-dimensional and 3-dimensional, but got ndim=" + + std::to_string(a_nd) + " and ndim=" + std::to_string(b_nd) + "."); + } + if (trans_op < 0 || trans_op > 2) { + throw py::value_error( + "trsm_batch: trans_op must be 0 (N), 1 (T), or 2 (C)."); + } + + const py::ssize_t *a_shape = matrixA.get_shape_raw(); + const py::ssize_t *b_shape = matrixB.get_shape_raw(); + + if (a_shape[0] != a_shape[1]) { + throw py::value_error("Expected batch of square matrices, but got " + "matrix shape (" + + std::to_string(a_shape[0]) + ", " + + std::to_string(a_shape[1]) + ") in batch"); + } + if (a_shape[0] != b_shape[0] || (a_nd == 3 && a_shape[2] != b_shape[2])) { + throw py::value_error("The shapes of the triangular and right-hand " + "sides arrays are incompatible."); + } + + // check compatibility of execution queue and allocation queue + if (!dpnp::utils::queues_are_compatible(exec_q, {matrixA, matrixB})) { + throw py::value_error( + "Execution queue is not compatible with allocation queues"); + } + + auto const &overlap = dpnp::tensor::overlap::MemoryOverlap(); + if (overlap(matrixA, matrixB)) { + throw py::value_error("The triangular and right-hand sides arrays " + "are overlapping segments of memory"); + } + + dpnp::tensor::validation::CheckWritable::throw_if_not_writable(matrixB); + + if (!matrixA.is_f_contiguous() || !matrixB.is_f_contiguous()) { + throw py::value_error("The triangular and right-hand sides arrays " + "must be F-contiguous"); + } + + auto array_types = dpnp_td_ns::usm_ndarray_types(); + const int a_type_id = + array_types.typenum_to_lookup_id(matrixA.get_typenum()); + const int b_type_id = + array_types.typenum_to_lookup_id(matrixB.get_typenum()); + + if (a_type_id != b_type_id) { + throw py::value_error("The types of the triangular and right-hand " + "sides arrays are mismatched"); + } + + trsm_batch_impl_fn_ptr_t trsm_batch_fn = + trsm_batch_dispatch_vector[a_type_id]; + if (trsm_batch_fn == nullptr) { + throw py::value_error( + "No trsm_batch implementation defined for the provided type " + "of the input matrix"); + } + + const std::int64_t m = a_shape[0]; + const std::int64_t n = b_shape[1]; + const std::int64_t batch_size = b_shape[2]; + + const std::int64_t lda = std::max(1, m); + const std::int64_t ldb = std::max(1, m); + // stride 0 reuses the shared matrix A for each batch element + const std::int64_t stride_a = (a_nd == 2) ? 0 : m * m; + const std::int64_t stride_b = m * n; + + const oneapi::mkl::uplo uplo_val = + static_cast(upper_lower); + const oneapi::mkl::transpose trans = + (trans_op == 0) ? oneapi::mkl::transpose::N + : (trans_op == 1) ? oneapi::mkl::transpose::T + : oneapi::mkl::transpose::C; + const oneapi::mkl::diag diag_val = + unit_diag ? oneapi::mkl::diag::unit : oneapi::mkl::diag::nonunit; + + sycl::event trsm_batch_ev = trsm_batch_fn( + exec_q, uplo_val, trans, diag_val, m, n, matrixA.get_data(), lda, + stride_a, matrixB.get_data(), ldb, stride_b, batch_size, depends); + + sycl::event args_ev = dpnp::utils::keep_args_alive( + exec_q, {matrixA, matrixB}, {trsm_batch_ev}); + + return std::make_pair(args_ev, trsm_batch_ev); +} + +template +struct TrsmBatchContigFactory +{ + fnT get() + { + if constexpr (types::TrsmBatchTypePairSupportFactory::is_defined) { + return trsm_batch_impl; + } + else { + return nullptr; + } + } +}; + +void init_trsm_batch_dispatch_vector(void) +{ + init_dispatch_vector( + trsm_batch_dispatch_vector); +} +} // namespace dpnp::extensions::blas diff --git a/dpnp/backend/extensions/blas/types_matrix.hpp b/dpnp/backend/extensions/blas/types_matrix.hpp index 5751b5634777..7bd53d180dc1 100644 --- a/dpnp/backend/extensions/blas/types_matrix.hpp +++ b/dpnp/backend/extensions/blas/types_matrix.hpp @@ -214,4 +214,30 @@ struct SyrkTypePairSupportFactory // fall-through dpnp_td_ns::NotDefinedEntry>::is_defined; }; + +/** + * @brief A factory to define pairs of supported types for which + * MKL BLAS library provides support in oneapi::mkl::blas::trsm_batch + * function. + * + * @tparam T Type of input triangular matrices and of the right-hand sides + * overwritten by the solutions. + */ +template +struct TrsmBatchTypePairSupportFactory +{ + static constexpr bool is_defined = + std::disjunction, + dpnp_td_ns::TypePairDefinedEntry, + dpnp_td_ns::TypePairDefinedEntry, + T, + std::complex>, + dpnp_td_ns::TypePairDefinedEntry, + T, + std::complex>, + // fall-through + dpnp_td_ns::NotDefinedEntry>::is_defined; +}; } // namespace dpnp::extensions::blas::types diff --git a/dpnp/backend/extensions/lapack/CMakeLists.txt b/dpnp/backend/extensions/lapack/CMakeLists.txt index ecb0c31d391f..3b33eaad9abc 100644 --- a/dpnp/backend/extensions/lapack/CMakeLists.txt +++ b/dpnp/backend/extensions/lapack/CMakeLists.txt @@ -47,6 +47,7 @@ set(_module_src ${CMAKE_CURRENT_SOURCE_DIR}/orgqr_batch.cpp ${CMAKE_CURRENT_SOURCE_DIR}/potrf.cpp ${CMAKE_CURRENT_SOURCE_DIR}/potrf_batch.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/potrs.cpp ${CMAKE_CURRENT_SOURCE_DIR}/syevd.cpp ${CMAKE_CURRENT_SOURCE_DIR}/syevd_batch.cpp ${CMAKE_CURRENT_SOURCE_DIR}/ungqr.cpp diff --git a/dpnp/backend/extensions/lapack/lapack_py.cpp b/dpnp/backend/extensions/lapack/lapack_py.cpp index eca276f719c4..ab2ec753db52 100644 --- a/dpnp/backend/extensions/lapack/lapack_py.cpp +++ b/dpnp/backend/extensions/lapack/lapack_py.cpp @@ -44,6 +44,7 @@ #include "linalg_exceptions.hpp" #include "orgqr.hpp" #include "potrf.hpp" +#include "potrs.hpp" #include "syevd.hpp" #include "syevd_batch.hpp" #include "ungqr.hpp" @@ -67,6 +68,7 @@ void init_dispatch_vectors(void) lapack_ext::init_orgqr_dispatch_vector(); lapack_ext::init_potrf_batch_dispatch_vector(); lapack_ext::init_potrf_dispatch_vector(); + lapack_ext::init_potrs_dispatch_vector(); lapack_ext::init_ungqr_batch_dispatch_vector(); lapack_ext::init_ungqr_dispatch_vector(); } @@ -214,6 +216,14 @@ PYBIND11_MODULE(_lapack_impl, m, py::mod_gil_not_used()) py::arg("n"), py::arg("stride_a"), py::arg("batch_size"), py::arg("depends") = py::list()); + m.def("_potrs", &lapack_ext::potrs, + "Call `potrs` from OneMKL LAPACK library to return " + "the solutions of linear equations with a Cholesky-factored " + "symmetric (Hermitian) positive-definite coefficient matrix, " + "with multiple right-hand sides", + py::arg("sycl_queue"), py::arg("a_array"), py::arg("b_array"), + py::arg("upper_lower"), py::arg("depends") = py::list()); + m.def("_ungqr_batch", &lapack_ext::ungqr_batch, "Call `_ungqr_batch` from OneMKL LAPACK library to return " "the complex unitary matrices matrix Qi of the QR factorization " diff --git a/dpnp/backend/extensions/lapack/potrs.cpp b/dpnp/backend/extensions/lapack/potrs.cpp new file mode 100644 index 000000000000..9775edf18ac4 --- /dev/null +++ b/dpnp/backend/extensions/lapack/potrs.cpp @@ -0,0 +1,280 @@ +//***************************************************************************** +// Copyright (c) 2026, Intel Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// - Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// - Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +// THE POSSIBILITY OF SUCH DAMAGE. +//***************************************************************************** + +#include + +#include + +// utils extension header +#include "ext/common.hpp" + +// dpnp tensor headers +#include "utils/memory_overlap.hpp" +#include "utils/output_validation.hpp" +#include "utils/sycl_alloc_utils.hpp" +#include "utils/type_utils.hpp" + +#include "linalg_exceptions.hpp" +#include "potrs.hpp" +#include "types_matrix.hpp" + +namespace dpnp::extensions::lapack +{ +namespace mkl_lapack = oneapi::mkl::lapack; +namespace py = pybind11; +namespace type_utils = dpnp::tensor::type_utils; + +using ext::common::init_dispatch_vector; + +typedef sycl::event (*potrs_impl_fn_ptr_t)(sycl::queue &, + const oneapi::mkl::uplo, + const std::int64_t, + const std::int64_t, + char *, + const std::int64_t, + char *, + const std::int64_t, + std::vector &, + const std::vector &); + +static potrs_impl_fn_ptr_t potrs_dispatch_vector[dpnp_td_ns::num_types]; + +template +static sycl::event potrs_impl(sycl::queue &exec_q, + const oneapi::mkl::uplo upper_lower, + const std::int64_t n, + const std::int64_t nrhs, + char *in_a, + const std::int64_t lda, + char *in_b, + const std::int64_t ldb, + std::vector &host_task_events, + const std::vector &depends) +{ + type_utils::validate_type_for_device(exec_q); + + T *a = reinterpret_cast(in_a); + T *b = reinterpret_cast(in_b); + + const std::int64_t scratchpad_size = mkl_lapack::potrs_scratchpad_size( + exec_q, upper_lower, n, nrhs, lda, ldb); + T *scratchpad = nullptr; + + std::stringstream error_msg; + std::int64_t info = 0; + bool is_exception_caught = false; + + sycl::event potrs_event; + try { + // Release GIL to avoid serialization of host task submissions + // to the same queue in OneMKL + py::gil_scoped_release lock{}; + + scratchpad = sycl::malloc_device(scratchpad_size, exec_q); + + potrs_event = mkl_lapack::potrs( + exec_q, + upper_lower, // An enumeration value of type oneapi::mkl::uplo: + // oneapi::mkl::uplo::upper if the Cholesky factor is + // stored in the upper triangular part of `a`; + // oneapi::mkl::uplo::lower if it is stored in the + // lower triangular part. + n, // The order of the matrix A (0 ≤ n). + nrhs, // The number of right-hand sides (0 ≤ nrhs). + a, // Pointer to the Cholesky-factored matrix (n x n). + lda, // The leading dimension of `a`. + b, // Pointer to the matrix B of right-hand sides (ldb, nrhs), + // overwritten by the solution matrix X. + ldb, // The leading dimension of `b`. + scratchpad, // Pointer to scratchpad memory to be used by MKL + // routine for storing intermediate results. + scratchpad_size, depends); + } catch (mkl_lapack::exception const &e) { + is_exception_caught = true; + info = e.info(); + + if (info < 0) { + error_msg << "Parameter number " << -info + << " had an illegal value."; + } + else if (info == scratchpad_size && e.detail() != 0) { + error_msg + << "Insufficient scratchpad size. Required size is at least " + << e.detail(); + } + else { + error_msg << "Unexpected MKL exception caught during potrs() " + "call:\nreason: " + << e.what() << "\ninfo: " << e.info(); + } + } catch (sycl::exception const &e) { + is_exception_caught = true; + error_msg << "Unexpected SYCL exception caught during potrs() call:\n" + << e.what(); + } + + if (is_exception_caught) // an unexpected error occurs + { + if (scratchpad != nullptr) { + dpnp::tensor::alloc_utils::sycl_free_noexcept(scratchpad, exec_q); + } + throw std::runtime_error(error_msg.str()); + } + + sycl::event clean_up_event = exec_q.submit([&](sycl::handler &cgh) { + cgh.depends_on(potrs_event); + auto ctx = exec_q.get_context(); + cgh.host_task([ctx, scratchpad]() { + dpnp::tensor::alloc_utils::sycl_free_noexcept(scratchpad, ctx); + }); + }); + host_task_events.push_back(clean_up_event); + return potrs_event; +} + +std::pair + potrs(sycl::queue &exec_q, + const dpnp::tensor::usm_ndarray &a_array, + const dpnp::tensor::usm_ndarray &b_array, + const std::int8_t upper_lower, + const std::vector &depends) +{ + const int a_array_nd = a_array.get_ndim(); + const int b_array_nd = b_array.get_ndim(); + + if (a_array_nd != 2) { + throw py::value_error("The Cholesky-factored array has ndim=" + + std::to_string(a_array_nd) + + ", but a 2-dimensional array is expected."); + } + if (b_array_nd > 2) { + throw py::value_error( + "The right-hand sides array has ndim=" + + std::to_string(b_array_nd) + + ", but a 1-dimensional or a 2-dimensional array is expected."); + } + + const py::ssize_t *a_array_shape = a_array.get_shape_raw(); + const py::ssize_t *b_array_shape = b_array.get_shape_raw(); + + if (a_array_shape[0] != a_array_shape[1]) { + throw py::value_error("The Cholesky-factored array must be square," + " but got a shape of (" + + std::to_string(a_array_shape[0]) + ", " + + std::to_string(a_array_shape[1]) + ")."); + } + if (a_array_shape[0] != b_array_shape[0]) { + throw py::value_error("The right-hand sides array must have " + + std::to_string(a_array_shape[0]) + + " rows, but got " + + std::to_string(b_array_shape[0]) + "."); + } + + // check compatibility of execution queue and allocation queue + if (!dpnp::utils::queues_are_compatible(exec_q, {a_array, b_array})) { + throw py::value_error( + "Execution queue is not compatible with allocation queues"); + } + + auto const &overlap = dpnp::tensor::overlap::MemoryOverlap(); + if (overlap(a_array, b_array)) { + throw py::value_error("The Cholesky-factored and right-hand sides " + "arrays are overlapping segments of memory"); + } + + dpnp::tensor::validation::CheckWritable::throw_if_not_writable(b_array); + + if (!a_array.is_f_contiguous()) { + throw py::value_error("The Cholesky-factored array " + "must be F-contiguous"); + } + if (!b_array.is_f_contiguous()) { + throw py::value_error("The right-hand sides array " + "must be F-contiguous"); + } + + auto array_types = dpnp_td_ns::usm_ndarray_types(); + int a_array_type_id = + array_types.typenum_to_lookup_id(a_array.get_typenum()); + int b_array_type_id = + array_types.typenum_to_lookup_id(b_array.get_typenum()); + + if (a_array_type_id != b_array_type_id) { + throw py::value_error("The types of the Cholesky-factored and " + "right-hand sides arrays are mismatched"); + } + + potrs_impl_fn_ptr_t potrs_fn = potrs_dispatch_vector[a_array_type_id]; + if (potrs_fn == nullptr) { + throw py::value_error( + "No potrs implementation defined for the provided type " + "of the input matrix."); + } + + const std::int64_t n = a_array_shape[0]; + const std::int64_t nrhs = (b_array_nd > 1) ? b_array_shape[1] : 1; + + const std::int64_t lda = std::max(1UL, n); + const std::int64_t ldb = std::max(1UL, n); + + char *a_array_data = a_array.get_data(); + char *b_array_data = b_array.get_data(); + const oneapi::mkl::uplo uplo_val = + static_cast(upper_lower); + + std::vector host_task_events; + sycl::event potrs_ev = + potrs_fn(exec_q, uplo_val, n, nrhs, a_array_data, lda, b_array_data, + ldb, host_task_events, depends); + + sycl::event args_ev = dpnp::utils::keep_args_alive( + exec_q, {a_array, b_array}, host_task_events); + + return std::make_pair(args_ev, potrs_ev); +} + +template +struct PotrsContigFactory +{ + fnT get() + { + if constexpr (types::PotrsTypePairSupportFactory::is_defined) { + return potrs_impl; + } + else { + return nullptr; + } + } +}; + +void init_potrs_dispatch_vector(void) +{ + init_dispatch_vector( + potrs_dispatch_vector); +} +} // namespace dpnp::extensions::lapack diff --git a/dpnp/backend/extensions/lapack/potrs.hpp b/dpnp/backend/extensions/lapack/potrs.hpp new file mode 100644 index 000000000000..f83996c63224 --- /dev/null +++ b/dpnp/backend/extensions/lapack/potrs.hpp @@ -0,0 +1,46 @@ +//***************************************************************************** +// Copyright (c) 2026, Intel Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// - Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// - Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +// THE POSSIBILITY OF SUCH DAMAGE. +//***************************************************************************** + +#pragma once + +#include +#include + +#include "dpnp4pybind11.hpp" + +namespace dpnp::extensions::lapack +{ +extern std::pair + potrs(sycl::queue &exec_q, + const dpnp::tensor::usm_ndarray &a_array, + const dpnp::tensor::usm_ndarray &b_array, + const std::int8_t upper_lower, + const std::vector &depends = {}); + +extern void init_potrs_dispatch_vector(void); +} // namespace dpnp::extensions::lapack diff --git a/dpnp/backend/extensions/lapack/types_matrix.hpp b/dpnp/backend/extensions/lapack/types_matrix.hpp index ea000bea5c88..e643e5547806 100644 --- a/dpnp/backend/extensions/lapack/types_matrix.hpp +++ b/dpnp/backend/extensions/lapack/types_matrix.hpp @@ -388,6 +388,34 @@ struct PotrfBatchTypePairSupportFactory dpnp_td_ns::NotDefinedEntry>::is_defined; }; +/** + * @brief A factory to define pairs of supported types for which + * MKL LAPACK library provides support in oneapi::mkl::lapack::potrs + * function. + * + * @tparam T Type of array containing input matrix (Cholesky-factored form) + * and the array of multiple dependent variables, + * as well as the output array for storing the solutions to a system of linear + * equations. + */ +template +struct PotrsTypePairSupportFactory +{ + static constexpr bool is_defined = + std::disjunction, + dpnp_td_ns::TypePairDefinedEntry, + dpnp_td_ns::TypePairDefinedEntry, + T, + std::complex>, + dpnp_td_ns::TypePairDefinedEntry, + T, + std::complex>, + // fall-through + dpnp_td_ns::NotDefinedEntry>::is_defined; +}; + /** * @brief A factory to define pairs of supported types for which * MKL LAPACK library provides support in oneapi::mkl::lapack::syevd diff --git a/dpnp/linalg/dpnp_iface_linalg.py b/dpnp/linalg/dpnp_iface_linalg.py index 8ed355278cbe..8b9dae4f90aa 100644 --- a/dpnp/linalg/dpnp_iface_linalg.py +++ b/dpnp/linalg/dpnp_iface_linalg.py @@ -131,6 +131,12 @@ def cholesky(a, /, *, upper=False): L : (..., M, M) dpnp.ndarray Lower or upper-triangular Cholesky factor of `a`. + See Also + -------- + :obj:`dpnp.scipy.linalg.cho_factor` : Cholesky decomposition of a matrix, + to use in + :obj:`dpnp.scipy.linalg.cho_solve`. + Examples -------- >>> import dpnp as np @@ -149,8 +155,8 @@ def cholesky(a, /, *, upper=False): The upper-triangular Cholesky factor can also be obtained: >>> np.linalg.cholesky(A, upper=True) - array([[ 1.+0.j, -0.-2.j], - [ 0.+0.j, 1.+0.j]] + array([[1., 2.], + [0., 1.]]) """ diff --git a/dpnp/scipy/linalg/__init__.py b/dpnp/scipy/linalg/__init__.py index 81eadd890fa9..01af1433e796 100644 --- a/dpnp/scipy/linalg/__init__.py +++ b/dpnp/scipy/linalg/__init__.py @@ -35,9 +35,12 @@ """ +from ._decomp_cholesky import cho_factor, cho_solve from ._decomp_lu import lu, lu_factor, lu_solve __all__ = [ + "cho_factor", + "cho_solve", "lu", "lu_factor", "lu_solve", diff --git a/dpnp/scipy/linalg/_decomp_cholesky.py b/dpnp/scipy/linalg/_decomp_cholesky.py new file mode 100644 index 000000000000..693b56b74ccc --- /dev/null +++ b/dpnp/scipy/linalg/_decomp_cholesky.py @@ -0,0 +1,239 @@ +# ***************************************************************************** +# Copyright (c) 2026, Intel Corporation +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# - Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# - Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# - Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE +# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +# THE POSSIBILITY OF SUCH DAMAGE. +# ***************************************************************************** + + +""" +Interface of the SciPy-compatible Linear Algebra subset for DPNP. + +Notes +----- +This module exposes the public API for ``dpnp.scipy.linalg``. +It contains: + - SciPy-like interface functions + - documentation for the functions + +""" + +import numpy + +import dpnp +from dpnp.linalg.dpnp_utils_linalg import assert_stacked_2d + +from ._utils import ( + dpnp_cho_factor, + dpnp_cho_solve, +) + + +def _lower_to_bool(lower): + """Convert `lower` to bool, also accepting an array of equal values.""" + + if dpnp.is_supported_array_type(lower): + lower = dpnp.asnumpy(lower) + if isinstance(lower, numpy.ndarray) and lower.size != 1: + if lower.size == 0: + return False + if not (lower == lower.flat[0]).all(): + raise ValueError("All values of `lower` must be equal.") + lower = lower.flat[0] + return bool(lower) + + +def cho_factor(a, lower=False, overwrite_a=False, check_finite=True): + """ + Compute the Cholesky decomposition of a matrix, to use in + :obj:`dpnp.scipy.linalg.cho_solve`. + + Returns a matrix containing the Cholesky decomposition, + ``A = L L*`` or ``A = U* U`` of a Hermitian positive-definite matrix `a`. + The return value can be directly used as the first parameter to + :obj:`dpnp.scipy.linalg.cho_solve`. + + For full documentation refer to :obj:`scipy.linalg.cho_factor`. + + Parameters + ---------- + a : (..., M, M) {dpnp.ndarray, usm_ndarray} + Matrix to be decomposed. + lower : bool, optional + Whether to compute the upper or lower triangular Cholesky + factorization. + + Default: ``False``. + overwrite_a : bool, optional + Whether to overwrite data in `a` (may improve performance). + + Default: ``False``. + check_finite : bool, optional + Whether to check that the input matrix contains only finite numbers. + Disabling may give a performance gain, but may result in problems + (crashes, non-termination) if the inputs do contain infinities or NaNs. + + Default: ``True``. + + Returns + ------- + c : (..., M, M) dpnp.ndarray + Matrix whose upper or lower triangle contains the Cholesky factor + of `a`. If `a` is not overwritten, the other triangle is zeroed, + otherwise it contains the data of `a` not used by the decomposition. + lower : bool + Flag indicating whether the factor is in the lower or upper triangle. + + Raises + ------ + dpnp.linalg.LinAlgError + Raised if decomposition fails. + + Warnings + -------- + This function synchronizes in order to validate array elements + when ``check_finite=True``, and to check the factorization status. + + See Also + -------- + :func:`dpnp.scipy.linalg.cho_solve` : Solve a linear set equations using + the Cholesky factorization of + a matrix. + :func:`dpnp.linalg.cholesky` : Cholesky decomposition of a matrix. + + Notes + ----- + During the finiteness check (if selected), the entire matrix `a` is + checked. During decomposition, `a` is assumed to be symmetric or Hermitian + (as applicable), and only the half selected by option `lower` is + referenced. + + Examples + -------- + >>> import dpnp as np + >>> A = np.array([[9., 3., 1., 5.], [3., 7., 5., 1.], + ... [1., 5., 9., 2.], [5., 1., 2., 6.]]) + >>> c, low = np.scipy.linalg.cho_factor(A) + >>> c + array([[ 3. , 1. , 0.33333333, 1.66666667], + [ 0. , 2.44948974, 1.90515869, -0.27216553], + [ 0. , 0. , 2.29330749, 0.8559528 ], + [ 0. , 0. , 0. , 1.55418563]]) + >>> np.allclose(c.T @ c, A) + array(True) + + Batched input: + + >>> B = np.stack([A, 2 * A]) + >>> c, low = np.scipy.linalg.cho_factor(B, lower=True) + >>> np.allclose(c @ c.mT, B) + array(True) + + """ + + dpnp.check_supported_arrays_type(a) + assert_stacked_2d(a) + + if a.shape[-1] != a.shape[-2]: + raise ValueError( + "Expected a square matrix or batch thereof, " + f"got a.shape={a.shape}" + ) + + return dpnp_cho_factor( + a, + lower=bool(lower), + overwrite_a=overwrite_a, + check_finite=check_finite, + ) + + +def cho_solve(c_and_lower, b, overwrite_b=False, check_finite=True): + """ + Solve the linear equations ``A x = b``, given the Cholesky factorization + of `A`. + + For full documentation refer to :obj:`scipy.linalg.cho_solve`. + + Parameters + ---------- + (c, lower) : tuple, ({dpnp.ndarray, usm_ndarray}, bool) + Cholesky factorization of `A` (..., M, M), as given by + :obj:`dpnp.scipy.linalg.cho_factor`. `lower` can also be an array + whose values are all equal. + b : {(M,), (..., M, K)} {dpnp.ndarray, usm_ndarray} + Right-hand side. + overwrite_b : bool, optional + Whether to overwrite data in `b` (may improve performance). + + Default: ``False``. + check_finite : bool, optional + Whether to check that the input matrices contain only finite numbers. + Disabling may give a performance gain, but may result in problems + (crashes, non-termination) if the inputs do contain infinities or NaNs. + + Default: ``True``. + + Returns + ------- + x : {(M,), (..., M, K)} dpnp.ndarray + The solution to the system ``A x = b``. + + Warnings + -------- + This function synchronizes in order to validate array elements + when ``check_finite=True``. + + See Also + -------- + :func:`dpnp.scipy.linalg.cho_factor` : Cholesky factorization of a matrix. + + Examples + -------- + >>> import dpnp as np + >>> A = np.array([[9., 3., 1., 5.], [3., 7., 5., 1.], + ... [1., 5., 9., 2.], [5., 1., 2., 6.]]) + >>> c, low = np.scipy.linalg.cho_factor(A) + >>> x = np.scipy.linalg.cho_solve((c, low), np.ones(4)) + >>> x + array([-0.01749271, 0.11953353, 0.01166181, 0.1574344 ]) + >>> np.allclose(A @ x, np.ones(4)) + array(True) + + """ + + c, lower = c_and_lower + dpnp.check_supported_arrays_type(c, b) + assert_stacked_2d(c) + + if c.shape[-1] != c.shape[-2]: + raise ValueError("The factored matrix c is not square.") + + return dpnp_cho_solve( + c, + _lower_to_bool(lower), + b, + overwrite_b=overwrite_b, + check_finite=check_finite, + ) diff --git a/dpnp/scipy/linalg/_utils.py b/dpnp/scipy/linalg/_utils.py index 6b0deb377211..d145b46a4138 100644 --- a/dpnp/scipy/linalg/_utils.py +++ b/dpnp/scipy/linalg/_utils.py @@ -46,6 +46,7 @@ import dpctl.utils as dpu import dpnp +import dpnp.backend.extensions.blas._blas_impl as bi import dpnp.backend.extensions.lapack._lapack_impl as li # pylint: disable=no-name-in-module @@ -54,22 +55,22 @@ from dpnp.linalg.dpnp_utils_linalg import _common_type, _real_type -def _align_lu_solve_broadcast(lu, b): - """Align LU and RHS batch dimensions with SciPy-like rules.""" +def _align_lu_solve_broadcast(lu, b, name="lu"): + """Align LU (or Cholesky) and RHS batch dimensions with SciPy-like rules.""" lu_shape = lu.shape b_shape = b.shape if b.ndim < 2: if lu_shape[-2] != b_shape[0]: raise ValueError( - f"Shapes of lu {lu_shape} and b {b_shape} are incompatible" + f"Shapes of {name} {lu_shape} and b {b_shape} are incompatible" ) b = dpnp.broadcast_to(b, lu_shape[:-1]) return lu, b if lu_shape[-2] != b_shape[-2]: raise ValueError( - f"Shapes of lu {lu_shape} and b {b_shape} are incompatible" + f"Shapes of {name} {lu_shape} and b {b_shape} are incompatible" ) # Use dpnp.broadcast_shapes() to align the resulting batch shapes @@ -351,18 +352,173 @@ def _batched_lu_solve(lu, piv, b, res_type, trans=0): return b_h -def _is_copy_required(a, res_type): +def _batched_cho_solve(c, lower, b, res_type, res_usm_type, exec_q): + """Solve a batched equation system given Cholesky factors.""" + + b_ndim_orig = b.ndim + c_2d = c if c.ndim == 2 else None + c, b = _align_lu_solve_broadcast(c, b, name="c") + + if b.size == 0: + return dpnp.empty( + b.shape, dtype=res_type, usm_type=res_usm_type, sycl_queue=exec_q + ) + + n = c.shape[-1] + nrhs = b.shape[-1] if b_ndim_orig > 1 else 1 + batch_size = c.size // (n * n) + + if c_2d is not None: + # a single factor is not broadcast: fold the batch into the RHS + # columns for one potrs call, except for a real type on GPU and + # a small batch, where oneMKL potrs slows down above 16 columns + # and the factor is shared by the batch in trsm_batch (stride 0) + c_f, uplo, conj = _cho_factor_f(c_2d, lower, res_type) + fold = not ( + exec_q.sycl_device.is_gpu + and not dpnp.issubdtype(res_type, dpnp.complexfloating) + and 16 < batch_size * nrhs + and batch_size <= 32 + ) + else: + fold = False + c = dpnp.reshape(c, (batch_size, n, n)) + c_is_complex = dpnp.issubdtype(c.dtype, dpnp.complexfloating) + c = dpnp.astype(c, res_type, copy=False) + + # oneMKL expects each factor in column-major order: (n, n, batch) + c_f = dpnp.transpose(c, (1, 2, 0)) + if c_f.flags.f_contiguous: + uplo, conj = int(lower), False + else: + # pass C-contiguous factors transposed, see `dpnp_cho_solve` + c = dpnp.ascontiguousarray(c) + c_f = dpnp.transpose(c, (2, 1, 0)) + uplo, conj = int(not lower), c_is_complex + + # the solves overwrite `b`: copy it to (n, nrhs, batch) layout + b_h = dpnp.empty( + (n, nrhs, batch_size), + dtype=res_type, + order="F", + usm_type=res_usm_type, + sycl_queue=exec_q, + ) + b_src = dpnp.reshape(b, (batch_size, n, nrhs)) + _copy_rhs(b_src, dpnp.moveaxis(b_h, -1, 0), conj, exec_q) + + if fold: + _potrs(c_f, dpnp.reshape(b_h, (n, -1), order="F"), uplo, exec_q) + else: + _cho_solve_trsm(c_f, b_h, uplo, exec_q) + + if conj: + dpnp.conj(b_h, out=b_h) + + # Restore original shape: move batch axis back and reshape + return dpnp.moveaxis(b_h, -1, 0).reshape(b.shape) + + +def _cho_solve_trsm(c_f, b_f, uplo, exec_q): + """ + Solve ``A X = B`` in place of `b_f` given the Cholesky factors `c_f`. + + `c_f` is a (n, n, batch) or a shared (n, n) and `b_f` a (n, nrhs, batch) + F-contiguous array. + Two batched triangular solves are used, since oneMKL potrs_batch is much + slower (see MKLD-20341). + + """ + + # A = F^H F for an upper factor F and A = F F^H for a lower one + # (oneapi::mkl::uplo: upper=0, lower=1; trans_op: 0=N, 2=C) + _manager = dpu.SequentialOrderManager[exec_q] + c_usm = dpnp.get_usm_ndarray(c_f) + for trans_op in (2, 0) if uplo == 0 else (0, 2): + # Call the BLAS extension function _trsm_batch + # to solve the systems with triangular matrices + ht_ev, trsm_ev = bi._trsm_batch( + exec_q, + c_usm, + dpnp.get_usm_ndarray(b_f), + uplo, + trans_op=trans_op, + depends=_manager.submitted_events, + ) + _manager.add_event_pair(ht_ev, trsm_ev) + + +def _cho_factor_f(c, lower, res_type): + """ + Return a 2-D factor `c` in the layout expected by oneMKL potrs and trsm. + + oneMKL is column-major: an F-contiguous `c` is used as is, otherwise + `c.T` is passed with flipped `uplo`, so the solves are with conj(A) + for a complex factor and the RHS and the solution have to be conjugated + (a real factor cast to a complex type needs none). + + """ + + c_is_complex = dpnp.issubdtype(c.dtype, dpnp.complexfloating) + c = dpnp.astype(c, res_type, copy=False) + if c.flags.f_contiguous: + return c, int(lower), False + return dpnp.ascontiguousarray(c).T, int(not lower), c_is_complex + + +def _potrs(c_f, b_f, uplo, exec_q): + """Solve ``A X = B`` in place of the F-contiguous `b_f` with potrs.""" + + # Call the LAPACK extension function _potrs + # to solve the system of linear equations with a Cholesky-factored + # coefficient matrix, with multiple right-hand sides. + _manager = dpu.SequentialOrderManager[exec_q] + ht_ev, potrs_ev = li._potrs( + exec_q, + dpnp.get_usm_ndarray(c_f), + dpnp.get_usm_ndarray(b_f), + uplo, + depends=_manager.submitted_events, + ) + _manager.add_event_pair(ht_ev, potrs_ev) + + +def _copy_rhs(src, dst, conj, exec_q): + """Copy `src` into `dst`, conjugating it if `conj` is ``True``.""" + + if conj and src.dtype == dst.dtype: + # fused copy and conjugation + dpnp.conj(src, out=dst) + return + + _manager = dpu.SequentialOrderManager[exec_q] + ht_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray( + src=dpnp.get_usm_ndarray(src), + dst=dst.get_array(), + sycl_queue=exec_q, + depends=_manager.submitted_events, + ) + _manager.add_event_pair(ht_ev, copy_ev) + + if conj and dpnp.iscomplexobj(src): + dpnp.conj(dst, out=dst) + + +def _is_copy_required(a, res_type, allow_c_contig=False): """ Determine if `a` needs to be copied before LU decomposition. This matches SciPy behavior: copy is needed unless input is suitable for in-place modification. + + With `allow_c_contig`, C-contiguous input is also suitable + (as for the Cholesky decomposition). """ if a.dtype != res_type: return True - if not a.flags["F_CONTIGUOUS"]: + if not (a.flags.f_contiguous or (allow_c_contig and a.flags.c_contiguous)): return True - if not a.flags["WRITABLE"]: + if not a.flags.writable: return True return False @@ -447,6 +603,150 @@ def _pivots_to_permutation(piv, m): return perm +def dpnp_cho_factor(a, lower=False, overwrite_a=False, check_finite=True): + """ + dpnp_cho_factor(a, lower=False, overwrite_a=False, check_finite=True) + + Compute the Cholesky decomposition (SciPy-compatible behavior). + + This function mimics the behavior of `scipy.linalg.cho_factor` including + support for `overwrite_a` and `check_finite`. Unless `a` is overwritten, + the triangle not used by the factorization is zeroed, as in SciPy 1.18. + + """ + + res_type = _common_type(a) + a_sycl_queue = a.sycl_queue + + if check_finite: + if not dpnp.isfinite(a).all(): + raise ValueError("array must not contain infs or NaNs") + + # accommodate empty arrays + if a.size == 0: + return dpnp.zeros_like(a, dtype=res_type, order="C"), lower + + _manager = dpu.SequentialOrderManager[a_sycl_queue] + + # SciPy-compatible behavior: factorize in place only 2D input + # of the result type which is writable and C- or F-contiguous + if ( + overwrite_a + and a.ndim == 2 + and not _is_copy_required(a, res_type, allow_c_contig=True) + ): + c = a + dep_ev = _manager.submitted_events + else: + a = dpnp.astype(a, res_type, copy=False) + + # potrf references one triangle only: copy it and zero the other + c = dpnp.empty_like(a, order="C") + tri_fn = ti._tril if lower else ti._triu + ht_ev, copy_ev = tri_fn( + src=dpnp.get_usm_ndarray(a), + dst=c.get_array(), + k=0, + sycl_queue=a_sycl_queue, + depends=_manager.submitted_events, + ) + _manager.add_event_pair(ht_ev, copy_ev) + dep_ev = [copy_ev] + + # oneMKL is column-major and sees a C-contiguous `c` transposed, + # so flip `uplo` for it (oneapi::mkl::uplo: upper=0, lower=1) + if c.ndim == 2: + c_usm = dpnp.get_usm_ndarray(c) + if c.flags.c_contiguous: + uplo = int(not lower) + else: + # F-contiguous `c` (overwrite_a=True): pass its C-contiguous view + c_usm, uplo = c_usm.T, int(lower) + + # Call the LAPACK extension function _potrf + # to compute the Cholesky decomposition + ht_ev, potrf_ev = li._potrf(a_sycl_queue, c_usm, uplo, depends=dep_ev) + else: + n = c.shape[-1] + batch_size = c.size // (n * n) + + # Call the LAPACK extension function _potrf_batch + # to compute the Cholesky decomposition of a batch of matrices + ht_ev, potrf_ev = li._potrf_batch( + a_sycl_queue, + dpnp.reshape(c, (batch_size, n, n)).get_array(), + int(not lower), + n, + n * n, + batch_size, + depends=dep_ev, + ) + _manager.add_event_pair(ht_ev, potrf_ev) + + return c, lower + + +def dpnp_cho_solve(c, lower, b, overwrite_b=False, check_finite=True): + """ + dpnp_cho_solve(c, lower, b, overwrite_b=False, check_finite=True) + + Solve an equation system given the Cholesky factorization + (SciPy-compatible behavior). + + This function mimics the behavior of `scipy.linalg.cho_solve` including + support for `overwrite_b` and `check_finite`. + + """ + + res_usm_type, exec_q = get_usm_allocations([c, b]) + res_type = _common_type(c, b) + + if check_finite: + # submit both checks before waiting for the first result + c_is_finite = dpnp.isfinite(c).all() + b_is_finite = dpnp.isfinite(b).all() + if not (c_is_finite and b_is_finite): + raise ValueError("array must not contain infs or NaNs") + + if c.ndim > 2 or b.ndim > 2: + # SciPy always copies each 2D slice, + # so `overwrite_b` is ignored here + return _batched_cho_solve(c, lower, b, res_type, res_usm_type, exec_q) + + if c.shape[1] != b.shape[0]: + raise ValueError(f"incompatible dimensions ({c.shape} and {b.shape})") + + # accommodate empty arrays + if b.size == 0: + return dpnp.empty_like(b, dtype=res_type, usm_type=res_usm_type) + + c_f, uplo, conj = _cho_factor_f(c, lower, res_type) + + # SciPy-compatible behavior + # Copy is required if: + # - overwrite_b is False (always copy), + # - dtype mismatch, + # - not F-contiguous, + # - not writeable + if not overwrite_b or _is_copy_required(b, res_type): + b_h = dpnp.empty_like( + b, order="F", dtype=res_type, usm_type=res_usm_type + ) + _copy_rhs(b, b_h, conj, exec_q) + else: + # input is suitable for in-place modification + b_h = b + if conj: + dpnp.conj(b_h, out=b_h) + + _potrs(c_f, b_h, uplo, exec_q) + + if conj: + dpnp.conj(b_h, out=b_h) + + return b_h + + def dpnp_lu_factor(a, overwrite_a=False, check_finite=True): """ dpnp_lu_factor(a, overwrite_a=False, check_finite=True) diff --git a/dpnp/tests/test_linalg.py b/dpnp/tests/test_linalg.py index 5504cad7675c..87c5e0609754 100644 --- a/dpnp/tests/test_linalg.py +++ b/dpnp/tests/test_linalg.py @@ -20,6 +20,7 @@ generate_random_numpy_array, get_all_dtypes, get_float_complex_dtypes, + get_float_dtypes, get_integer_dtypes, get_integer_float_dtypes, has_support_aspect64, @@ -143,6 +144,405 @@ def test_usm_ndarray_linearalgebra_batch(func): assert isinstance(result, dpnp.ndarray) +def _make_hpd_np(shape, dtype, order="C"): + """Random Hermitian positive-definite matrix (or a batch of them).""" + a = generate_random_numpy_array(shape, dtype, low=-1, high=1) + n = shape[-1] + a = a @ numpy.swapaxes(a, -1, -2).conj() + n * numpy.eye(n, dtype=dtype) + return numpy.array(a, order=order) + + +def _tri(a, lower): + xp = numpy if isinstance(a, numpy.ndarray) else dpnp + return xp.tril(a) if lower else xp.triu(a) + + +def _other_tri(a, lower): + xp = numpy if isinstance(a, numpy.ndarray) else dpnp + return xp.triu(a, k=1) if lower else xp.tril(a, k=-1) + + +def _cholesky_np(a, lower): + l_np = numpy.linalg.cholesky(a) + return l_np if lower else numpy.swapaxes(l_np, -1, -2).conj() + + +class TestChoFactor: + @pytest.mark.parametrize("shape", [(1, 1), (2, 2), (3, 3), (5, 5)]) + @pytest.mark.parametrize("lower", [True, False]) + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("dtype", ALL_DTYPES_NO_BOOL) + def test_cho_factor(self, shape, lower, order, dtype): + a_np = _make_hpd_np(shape, dtype, order) + a_dp = dpnp.array(a_np, order=order) + + c, low = dpnp.scipy.linalg.cho_factor(a_dp, lower=lower) + assert low is lower + assert c.flags.c_contiguous + + expected = _cholesky_np(a_np, lower) + assert_dtype_allclose( + _tri(c, lower), expected, check_only_type_kind=True + ) + assert (_other_tri(c, lower) == 0).all() + + @pytest.mark.parametrize("shape", [(1, 1, 1), (2, 3, 3), (3, 2, 4, 4)]) + @pytest.mark.parametrize("lower", [True, False]) + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("dtype", get_float_complex_dtypes()) + def test_batched(self, shape, lower, order, dtype): + a_np = _make_hpd_np(shape, dtype, order) + a_dp = dpnp.array(a_np, order=order) + + c, low = dpnp.scipy.linalg.cho_factor(a_dp, lower=lower) + assert low is lower + assert c.shape == shape + + expected = _cholesky_np(a_np, lower) + assert_dtype_allclose(_tri(c, lower), expected) + assert (_other_tri(c, lower) == 0).all() + + @pytest.mark.parametrize("lower", [True, False]) + @pytest.mark.parametrize("dtype", get_float_complex_dtypes()) + def test_unused_triangle_ignored(self, lower, dtype): + a_np = _make_hpd_np((4, 4), dtype) + expected = _cholesky_np(a_np, lower) + + # only the selected triangle has to be referenced + a_np[ + numpy.triu_indices(4, 1) if lower else numpy.tril_indices(4, -1) + ] = numpy.nan + a_dp = dpnp.array(a_np) + + c, _ = dpnp.scipy.linalg.cho_factor( + a_dp, lower=lower, check_finite=False + ) + assert_dtype_allclose(_tri(c, lower), expected) + + @pytest.mark.parametrize("lower", [True, False]) + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("dtype", get_float_complex_dtypes()) + def test_overwrite_inplace(self, lower, order, dtype): + a_np = _make_hpd_np((4, 4), dtype, order) + a_dp = dpnp.array(a_np, order=order) + + c, _ = dpnp.scipy.linalg.cho_factor(a_dp, lower=lower, overwrite_a=True) + assert c is a_dp + + assert_dtype_allclose(_tri(c, lower), _cholesky_np(a_np, lower)) + # the other triangle keeps the input data + assert_array_equal(_other_tri(c, lower), _other_tri(a_np, lower)) + + @pytest.mark.parametrize("case", ["int", "batched", "strided"]) + def test_overwrite_copy(self, case): + if case == "int": + a_dp = dpnp.array([[4, 2], [2, 3]], dtype=dpnp.int32) + elif case == "batched": + a_dp = dpnp.ones((2, 2, 2)) + dpnp.eye(2) + else: + a_dp = (dpnp.ones((4, 4)) + 4 * dpnp.eye(4))[::2, ::2] + a_orig = a_dp.copy() + + c, _ = dpnp.scipy.linalg.cho_factor(a_dp, overwrite_a=True) + assert c is not a_dp + assert_array_equal(a_dp, a_orig) + + def test_overwrite_read_only(self): + a_dp = dpnp.array([[4.0, 2.0], [2.0, 3.0]]) + a_dp.flags["W"] = False + a_orig = a_dp.copy() + + c, _ = dpnp.scipy.linalg.cho_factor(a_dp, overwrite_a=True) + assert c is not a_dp + assert_array_equal(a_dp, a_orig) + + @pytest.mark.parametrize("shape", [(0, 0), (2, 0, 0), (0, 2, 2)]) + def test_empty(self, shape): + a_dp = dpnp.empty(shape) + c, low = dpnp.scipy.linalg.cho_factor(a_dp) + assert c.shape == shape + assert c.dtype == dpnp.default_float_type(sycl_queue=a_dp.sycl_queue) + assert low is False + + def test_not_positive_definite(self): + a_dp = dpnp.array([[1.0, 2.0], [2.0, 1.0]]) + assert_raises( + dpnp.linalg.LinAlgError, dpnp.scipy.linalg.cho_factor, a_dp + ) + + # failing matrix first: on a WSL Arrow Lake iGPU, single-precision + # potrf_batch was seen to miss failures at odd batch indices + a_dp = dpnp.stack([-dpnp.eye(3), dpnp.eye(3), dpnp.eye(3)]) + assert_raises( + dpnp.linalg.LinAlgError, dpnp.scipy.linalg.cho_factor, a_dp + ) + + @pytest.mark.parametrize("bad", [numpy.inf, -numpy.inf, numpy.nan]) + def test_check_finite_raises(self, bad): + a_dp = dpnp.array([[4.0, 2.0], [2.0, bad]]) + assert_raises(ValueError, dpnp.scipy.linalg.cho_factor, a_dp) + + def test_errors(self): + a_dp = dpnp.ones((2, 3)) + assert_raises(ValueError, dpnp.scipy.linalg.cho_factor, a_dp) + + a_dp = dpnp.ones(3) + assert_raises( + dpnp.linalg.LinAlgError, dpnp.scipy.linalg.cho_factor, a_dp + ) + + a_np = numpy.eye(2) + assert_raises(TypeError, dpnp.scipy.linalg.cho_factor, a_np) + + +class TestChoSolve: + @pytest.mark.parametrize("n", [1, 2, 3, 5]) + @pytest.mark.parametrize("rhs_cols", [None, 1, 3]) + @pytest.mark.parametrize("lower", [True, False]) + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("dtype", ALL_DTYPES_NO_BOOL) + def test_cho_solve(self, n, rhs_cols, lower, order, dtype): + a_np = _make_hpd_np((n, n), dtype) + b_shape = (n,) if rhs_cols is None else (n, rhs_cols) + b_np = generate_random_numpy_array(b_shape, dtype, order) + + c, low = dpnp.scipy.linalg.cho_factor(dpnp.array(a_np), lower=lower) + # exercise both C- and F-contiguous factors + c = dpnp.array(c, order=order) + b_dp = dpnp.array(b_np, order=order) + + x = dpnp.scipy.linalg.cho_solve((c, low), b_dp) + assert x.shape == b_shape + assert_array_equal(b_dp, b_np) + + expected = numpy.linalg.solve(a_np, b_np) + assert_dtype_allclose(x, expected, check_only_type_kind=True) + + @pytest.mark.parametrize("lower", [True, False]) + @pytest.mark.parametrize("dtype", get_float_complex_dtypes()) + def test_cholesky_factor(self, lower, dtype): + a_np = _make_hpd_np((4, 4), dtype) + b_np = generate_random_numpy_array((4, 2), dtype) + + # a factor with zeroed other triangle, e.g. from dpnp.linalg.cholesky + c = dpnp.linalg.cholesky(dpnp.array(a_np), upper=not lower) + x = dpnp.scipy.linalg.cho_solve((c, lower), dpnp.array(b_np)) + + assert_dtype_allclose(x, numpy.linalg.solve(a_np, b_np)) + + @pytest.mark.parametrize("c_dtype", get_float_complex_dtypes()) + @pytest.mark.parametrize("b_dtype", ALL_DTYPES_NO_BOOL) + def test_mixed_dtypes(self, c_dtype, b_dtype): + a_np = _make_hpd_np((3, 3), c_dtype) + b_np = generate_random_numpy_array((3, 2), b_dtype) + + c, low = dpnp.scipy.linalg.cho_factor(dpnp.array(a_np)) + x = dpnp.scipy.linalg.cho_solve((c, low), dpnp.array(b_np)) + + # the accuracy is limited by the precision of the factor + tol = 10 * numpy.finfo(c_dtype).resolution + expected = numpy.linalg.solve(a_np, b_np) + assert_allclose(x, expected, rtol=tol, atol=tol) + + @pytest.mark.parametrize("c_shape", [(4, 4), (2, 4, 4)]) + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("overwrite_b", [True, False]) + @pytest.mark.parametrize("dtype", get_float_dtypes()) + def test_real_factor_complex_rhs(self, c_shape, order, overwrite_b, dtype): + # a real factor needs no conjugation even for a complex result + b_dtype = numpy.result_type(dtype, numpy.complex64) + a_np = _make_hpd_np(c_shape, dtype) + b_np = generate_random_numpy_array(c_shape[:-1] + (2,), b_dtype) + + c, low = dpnp.scipy.linalg.cho_factor(dpnp.array(a_np)) + c = dpnp.array(c, order=order) + b_dp = dpnp.array(b_np, order="F") + + x = dpnp.scipy.linalg.cho_solve((c, low), b_dp, overwrite_b=overwrite_b) + if overwrite_b and len(c_shape) == 2: + assert x is b_dp + assert_dtype_allclose(x, numpy.linalg.solve(a_np, b_np)) + + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("dtype", get_float_complex_dtypes()) + def test_overwrite_inplace(self, order, dtype): + a_np = _make_hpd_np((3, 3), dtype) + b_np = generate_random_numpy_array((3, 2), dtype) + + c, low = dpnp.scipy.linalg.cho_factor(dpnp.array(a_np)) + c = dpnp.array(c, order=order) + b_dp = dpnp.array(b_np, order="F") + + x = dpnp.scipy.linalg.cho_solve((c, low), b_dp, overwrite_b=True) + assert x is b_dp + assert_dtype_allclose(x, numpy.linalg.solve(a_np, b_np)) + + @pytest.mark.parametrize("case", ["C-contig", "int", "strided"]) + def test_overwrite_copy(self, case): + if case == "C-contig": + b_dp = dpnp.ones((3, 2), order="C") + elif case == "int": + b_dp = dpnp.ones((3, 2), dtype=dpnp.int32, order="F") + else: + b_dp = dpnp.ones((6, 2), order="F")[::2] + + c, low = dpnp.scipy.linalg.cho_factor(4 * dpnp.eye(3)) + b_orig = b_dp.copy() + + x = dpnp.scipy.linalg.cho_solve((c, low), b_dp, overwrite_b=True) + assert x is not b_dp + assert_array_equal(b_dp, b_orig) + assert_allclose(x, b_orig / 4) + + def test_overwrite_read_only(self): + c, low = dpnp.scipy.linalg.cho_factor(4 * dpnp.eye(2)) + b_dp = dpnp.ones(2) + b_dp.flags["W"] = False + + x = dpnp.scipy.linalg.cho_solve((c, low), b_dp, overwrite_b=True) + assert x is not b_dp + assert_allclose(x, numpy.full(2, 0.25)) + + @pytest.mark.parametrize( + "c_shape, b_shape", + [ + ((1, 2, 2), (2,)), + ((2, 4, 4), (4,)), + ((2, 4, 4), (4, 3)), + ((2, 4, 4), (2, 4, 1)), + ((2, 4, 4), (1, 4, 3)), + ((2, 3, 4, 4), (1, 3, 4, 2)), + ((2, 3, 4, 4), (2, 1, 4, 2)), + ((3, 4, 4), (1, 1, 4, 2)), + ((4, 4), (2, 4, 3)), + ((4, 4), (2, 3, 4, 1)), + ], + ) + @pytest.mark.parametrize("lower", [True, False]) + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("dtype", get_float_complex_dtypes()) + def test_batched(self, c_shape, b_shape, lower, order, dtype): + a_np = _make_hpd_np(c_shape, dtype) + b_np = generate_random_numpy_array(b_shape, dtype, order) + + c, low = dpnp.scipy.linalg.cho_factor(dpnp.array(a_np), lower=lower) + c = dpnp.array(c, order=order) + b_dp = dpnp.array(b_np, order=order) + + x = dpnp.scipy.linalg.cho_solve((c, low), b_dp) + + if b_np.ndim == 1: + b_exp = numpy.broadcast_to(b_np, c_shape[:-1])[..., None] + expected = numpy.linalg.solve(a_np, b_exp)[..., 0] + else: + expected = numpy.linalg.solve(a_np, b_np) + assert x.shape == expected.shape + assert_dtype_allclose(x, expected) + assert_array_equal(b_dp, b_np) + + # few columns are folded into one solve, many use a shared factor + @pytest.mark.parametrize("b_shape", [(2, 3, 4, 2), (3, 4, 8)]) + @pytest.mark.parametrize("order", ["C", "F"]) + @pytest.mark.parametrize("c_dtype", get_float_complex_dtypes()) + @pytest.mark.parametrize("b_dtype", get_float_complex_dtypes()) + def test_2d_factor_batched_rhs(self, b_shape, order, c_dtype, b_dtype): + # a 2-D factor is not broadcast to the RHS batch + a_np = _make_hpd_np((4, 4), c_dtype) + b_np = generate_random_numpy_array(b_shape, b_dtype) + + c, low = dpnp.scipy.linalg.cho_factor(dpnp.array(a_np)) + c = dpnp.array(c, order=order) + b_dp = dpnp.array(b_np) + + x = dpnp.scipy.linalg.cho_solve((c, low), b_dp) + assert x.shape == b_np.shape + assert_array_equal(b_dp, b_np) + + tol = 10 * numpy.finfo(c_dtype).resolution + expected = numpy.linalg.solve(a_np, b_np) + assert_allclose(x, expected, rtol=tol, atol=tol) + + @pytest.mark.parametrize("lower", [True, False]) + @pytest.mark.parametrize("xp", [numpy, dpnp]) + def test_lower_array(self, lower, xp): + # SciPy >= 1.18 returns `lower` tiled to the batch shape + dtype = dpnp.default_float_type() + a_np = _make_hpd_np((3, 4, 4), dtype) + b_np = generate_random_numpy_array((3, 4, 2), dtype) + + c, _ = dpnp.scipy.linalg.cho_factor(dpnp.array(a_np), lower=lower) + expected = numpy.linalg.solve(a_np, b_np) + for low in [ + xp.asarray(lower), + xp.full(3, lower), + xp.full((3, 1), lower), + ]: + x = dpnp.scipy.linalg.cho_solve((c, low), dpnp.array(b_np)) + assert_dtype_allclose(x, expected) + + def test_lower_array_errors(self): + c, _ = dpnp.scipy.linalg.cho_factor(4 * dpnp.ones((2, 1, 1))) + low = numpy.array([True, False]) + assert_raises( + ValueError, dpnp.scipy.linalg.cho_solve, (c, low), dpnp.ones(1) + ) + + @pytest.mark.parametrize( + "c_shape, b_shape, x_shape", + [ + ((0, 0), (0,), (0,)), + ((0, 0), (0, 2), (0, 2)), + ((2, 2), (2, 0), (2, 0)), + ((2, 0, 0), (0,), (2, 0)), + ((0, 2, 2), (2,), (0, 2)), + ((2, 2, 2), (2, 2, 0), (2, 2, 0)), + ], + ) + def test_empty(self, c_shape, b_shape, x_shape): + c = dpnp.empty(c_shape) + b = dpnp.empty(b_shape) + + x = dpnp.scipy.linalg.cho_solve((c, False), b) + assert x.shape == x_shape + assert x.dtype == dpnp.default_float_type(sycl_queue=c.sycl_queue) + + @pytest.mark.parametrize("bad", [numpy.inf, -numpy.inf, numpy.nan]) + def test_check_finite_raises(self, bad): + c, low = dpnp.scipy.linalg.cho_factor(4 * dpnp.eye(2)) + b_dp = dpnp.array([1.0, bad]) + assert_raises(ValueError, dpnp.scipy.linalg.cho_solve, (c, low), b_dp) + + # unlike CuPy and like SciPy, the whole `c` is checked + c[1, 0] = bad + assert_raises( + ValueError, dpnp.scipy.linalg.cho_solve, (c, low), dpnp.ones(2) + ) + + def test_errors(self): + c = dpnp.eye(3) + assert_raises( + ValueError, dpnp.scipy.linalg.cho_solve, (c, False), dpnp.ones(2) + ) + assert_raises( + ValueError, + dpnp.scipy.linalg.cho_solve, + (dpnp.ones((2, 3, 3)), False), + dpnp.ones((2, 2, 1)), + ) + assert_raises( + ValueError, + dpnp.scipy.linalg.cho_solve, + (dpnp.ones((3, 2)), False), + dpnp.ones(3), + ) + assert_raises( + TypeError, + dpnp.scipy.linalg.cho_solve, + (numpy.eye(2), False), + dpnp.ones(2), + ) + + class TestCholesky: @pytest.mark.parametrize( "array", diff --git a/dpnp/tests/test_sycl_queue.py b/dpnp/tests/test_sycl_queue.py index e0f882f2175d..3d212f602575 100644 --- a/dpnp/tests/test_sycl_queue.py +++ b/dpnp/tests/test_sycl_queue.py @@ -1739,6 +1739,39 @@ def test_lu_solve(self, a_data, b_data, device): assert_sycl_queue_equal(result.sycl_queue, a.sycl_queue) assert_sycl_queue_equal(result.sycl_queue, b.sycl_queue) + @pytest.mark.parametrize( + "data", + [ + [[4.0, 2.0], [2.0, 3.0]], + numpy.empty((0, 0)), + [[[4.0, 2.0], [2.0, 3.0]]], + numpy.empty((1, 0, 0)), + ], + ) + def test_cho_factor(self, data, device): + a = dpnp.array(data, device=device) + c, _ = dpnp.scipy.linalg.cho_factor(a) + assert_sycl_queue_equal(c.sycl_queue, a.sycl_queue) + + @pytest.mark.parametrize( + "a_data, b_data", + [ + ([[4.0, 2.0], [2.0, 3.0]], [1.0, 2.0]), + ([[4.0, 2.0], [2.0, 3.0]], numpy.empty((2, 0))), + ([[[4.0, 2.0], [2.0, 3.0]]], [1.0, 2.0]), + ([[[4.0, 2.0], [2.0, 3.0]]], numpy.empty((1, 2, 0))), + ], + ) + def test_cho_solve(self, a_data, b_data, device): + a = dpnp.array(a_data, device=device) + c, low = dpnp.scipy.linalg.cho_factor(a) + b = dpnp.array(b_data, device=device) + + result = dpnp.scipy.linalg.cho_solve((c, low), b) + + assert_sycl_queue_equal(result.sycl_queue, a.sycl_queue) + assert_sycl_queue_equal(result.sycl_queue, b.sycl_queue) + @pytest.mark.parametrize("n", [-1, 0, 1, 2, 3]) def test_matrix_power(self, n, device): x = dpnp.array([[1.0, 2.0], [3.0, 5.0]], device=device) diff --git a/dpnp/tests/test_usm_type.py b/dpnp/tests/test_usm_type.py index 567078c3eff2..7ae08f7eef78 100644 --- a/dpnp/tests/test_usm_type.py +++ b/dpnp/tests/test_usm_type.py @@ -1597,6 +1597,45 @@ def test_lu_solve(self, a_data, b_data, usm_type, usm_type_rhs): [usm_type, usm_type_rhs] ) + @pytest.mark.parametrize( + "data", + [ + [[4.0, 2.0], [2.0, 3.0]], + numpy.empty((0, 0)), + [[[4.0, 2.0], [2.0, 3.0]]], + numpy.empty((1, 0, 0)), + ], + ) + def test_cho_factor(self, data, usm_type): + a = dpnp.array(data, usm_type=usm_type) + c, _ = dpnp.scipy.linalg.cho_factor(a) + + assert a.usm_type == usm_type + assert c.usm_type == usm_type + + @pytest.mark.parametrize("usm_type_rhs", list_of_usm_types) + @pytest.mark.parametrize( + "a_data, b_data", + [ + ([[4.0, 2.0], [2.0, 3.0]], [1.0, 2.0]), + ([[4.0, 2.0], [2.0, 3.0]], numpy.empty((2, 0))), + ([[[4.0, 2.0], [2.0, 3.0]]], [1.0, 2.0]), + ([[[4.0, 2.0], [2.0, 3.0]]], numpy.empty((1, 2, 0))), + ], + ) + def test_cho_solve(self, a_data, b_data, usm_type, usm_type_rhs): + a = dpnp.array(a_data, usm_type=usm_type) + c, low = dpnp.scipy.linalg.cho_factor(a) + b = dpnp.array(b_data, usm_type=usm_type_rhs) + + result = dpnp.scipy.linalg.cho_solve((c, low), b) + + assert c.usm_type == usm_type + assert b.usm_type == usm_type_rhs + assert result.usm_type == dpt.get_coerced_usm_type( + [usm_type, usm_type_rhs] + ) + @pytest.mark.parametrize("n", [-1, 0, 1, 2, 3]) def test_matrix_power(self, n, usm_type): a = dpnp.array([[1, 2], [3, 5]], usm_type=usm_type) diff --git a/dpnp/tests/third_party/cupyx/linalg_tests/__init__.py b/dpnp/tests/third_party/cupyx/linalg_tests/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/dpnp/tests/third_party/cupyx/linalg_tests/test_solve.py b/dpnp/tests/third_party/cupyx/linalg_tests/test_solve.py new file mode 100644 index 000000000000..baddca0c406d --- /dev/null +++ b/dpnp/tests/third_party/cupyx/linalg_tests/test_solve.py @@ -0,0 +1,221 @@ +from __future__ import annotations + +import unittest + +import numpy +import pytest + +import dpnp as cupy +import dpnp as cupyx +from dpnp.tests.helper import has_support_aspect64 +from dpnp.tests.third_party.cupy import testing + +# from cupyx import cusolver + + +@testing.parameterize( + *testing.product( + { + "size": [5, 9, 17, 33], + "dtype": [ + numpy.float32, + numpy.float64, + numpy.complex64, + numpy.complex128, + ], + } + ) +) +@pytest.mark.skip("invh is not supported") +class TestInvh(unittest.TestCase): + + @testing.numpy_cupy_allclose(atol=1e-5) + def test_invh(self, xp): + a = self._create_symmetric_matrix(xp, self.size, self.dtype) + if xp == cupy: + return cupyx.linalg.invh(a) + else: + return numpy.linalg.inv(a) + + def _create_symmetric_matrix(self, xp, n, dtype): + if dtype == numpy.complex128: + f_dtype = numpy.float64 + elif dtype == numpy.complex64: + f_dtype = numpy.float32 + else: + f_dtype = dtype + a = testing.shaped_random((n, n), xp, f_dtype, scale=1) + a = a + a.T + xp.eye(n, dtype=f_dtype) * n + if dtype in (numpy.complex64, numpy.complex128): + b = testing.shaped_random((n, n), xp, f_dtype, scale=1) + b = b - b.T + a = a + 1j * b + return a + + +@testing.parameterize( + *testing.product( + { + "size": [8], + "dtype": [ + numpy.float32, + numpy.float64, + numpy.complex64, + numpy.complex128, + ], + } + ) +) +@pytest.mark.skip("invh is not supported") +class TestErrorInvh(unittest.TestCase): + + # @pytest.mark.skipif( + # cupy.cuda.runtime.runtimeGetVersion() == 12000, + # reason='This fails with CUDA 12.0.0 but pass in CUDA 12.0.1. (#7309)') + def test_invh(self): + a = self._create_symmetric_matrix(self.size, self.dtype) + with cupyx.errstate(linalg="raise"): + with self.assertRaises(numpy.linalg.LinAlgError): + cupyx.linalg.invh(a) + + def _create_symmetric_matrix(self, n, dtype): + a = testing.shaped_random((n, n), cupy, dtype, scale=1) + a = a + a.T - cupy.eye(n, dtype=dtype) + return a + + +# TODO: cusolver does not support nrhs > 1 for potrsBatched +@testing.parameterize( + *testing.product( + { + "shape": [(2, 3, 3)], + "dtype": [ + numpy.float32, + numpy.float64, + numpy.complex64, + numpy.complex128, + ], + } + ) +) +@pytest.mark.skip("invh is not supported") +class TestXFailBatchedInvh(unittest.TestCase): + + def test_invh(self): + # if not cusolver.check_availability('potrsBatched'): + # pytest.skip('potrsBatched is not available') + a = self._create_symmetric_matrix(self.shape, self.dtype) + with cupyx.errstate(linalg="raise"): + with self.assertRaises(numpy.linalg.LinAlgError): + cupyx.linalg.invh(a) + + def _create_symmetric_matrix(self, shape, dtype): + a = testing.shaped_random(shape, cupy, dtype, scale=1) + a = a @ a.transpose(0, 2, 1) + return a + + +@testing.parameterize( + *testing.product( + { + "shape": [(3, 4, 2, 2), (5, 3, 3), (7, 7)], + "nrhs": [None, 1, 8], + "lower": [True, False], + "order": ["C", "F"], + "overwrite_b": [True, False], + } + ) +) +class TestChosolve(unittest.TestCase): + + @staticmethod + def _solve(a, b): + if ( + numpy.lib.NumpyVersion(numpy.__version__) < "2.0.0" + or a.shape[:-1] != b.shape + ): + return numpy.linalg.solve(a, b) + b = b[..., numpy.newaxis] + return numpy.linalg.solve(a, b)[..., 0] + + @staticmethod + def _cho_solve(c_and_lower, b, overwrite_b): + # dpnp follows SciPy: a batched right-hand side is (..., n, k), so + # pass a batch of vectors with a trailing axis of size one + c = c_and_lower[0] + if c.ndim > 2 and b.ndim == c.ndim - 1: + return cupyx.scipy.linalg.cho_solve( + c_and_lower, + b[..., None], + overwrite_b=overwrite_b, + check_finite=False, + )[..., 0] + + # dpnp checks the whole factor for infs and NaNs, as SciPy does, + # so skip the check for NaNs in the unused triangle + return cupyx.scipy.linalg.cho_solve( + c_and_lower, b, overwrite_b=overwrite_b, check_finite=False + ) + + @testing.for_dtypes("fdFD") + @testing.numpy_cupy_allclose( + atol=1e-5, contiguous_check=False, type_check=has_support_aspect64() + ) + def test_chosolve(self, xp, dtype): + + # if (len(self.shape) > 2 and + # not cusolver.check_availability('potrsBatched')): + # pytest.skip('potrsBatched is not available') + + a = self._create_posdef_matrix(xp, self.shape, dtype) + b_shape = list(self.shape[:-1]) + if self.nrhs is not None: + b_shape.append(self.nrhs) + b = testing.shaped_random(b_shape, xp, dtype=dtype) + + if xp == cupy: + L = xp.linalg.cholesky(a) + if self.lower: + indexes = numpy.triu_indices(L.shape[-1], 1) + else: + L = xp.moveaxis(L, -1, -2).conj() + indexes = numpy.tril_indices(L.shape[-1], -1) + L[..., indexes[0], indexes[1]] = numpy.nan + L = cupy.asarray(L, order=self.order) + return self._cho_solve( + (L, self.lower), b, overwrite_b=self.overwrite_b + ) + else: + return self._solve(a, b) + + @testing.for_dtypes("fdFD") + def test_chosolve_does_not_overwrite_b(self, dtype): + + # if (len(self.shape) > 2 and + # not cusolver.check_availability('potrsBatched')): + # pytest.skip('potrsBatched is not available') + + a = self._create_posdef_matrix(cupy, self.shape, dtype) + b_shape = list(self.shape[:-1]) + if self.nrhs is not None: + b_shape.append(self.nrhs) + b = testing.shaped_random(b_shape, cupy, dtype=dtype) + + L = cupy.linalg.cholesky(a) + if self.lower: + indexes = numpy.triu_indices(L.shape[-1], 1) + else: + L = cupy.moveaxis(L, -1, -2).conj() + indexes = numpy.tril_indices(L.shape[-1], -1) + L[..., indexes[0], indexes[1]] = numpy.nan + L = cupy.asarray(L, order=self.order) + b_copy = cupy.asarray(b, copy=True) + self._cho_solve((L, self.lower), b, overwrite_b=False) + assert cupy.all(b == b_copy) + + def _create_posdef_matrix(self, xp, shape, dtype): + n = shape[-1] + a = testing.shaped_random(shape, xp, dtype, scale=1) + a = a @ a.swapaxes(-2, -1).conjugate() + a = a + n * xp.eye(n) + return a