diff --git a/.github/workflows/conda-package.yml b/.github/workflows/conda-package.yml index 1a6650798e92..9b6b093801ac 100644 --- a/.github/workflows/conda-package.yml +++ b/.github/workflows/conda-package.yml @@ -31,6 +31,7 @@ env: test_usm_type.py third_party/cupy/core_tests third_party/cupy/linalg_tests/test_product.py + third_party/cupy/linalg_tests/test_solve.py third_party/cupy/logic_tests/test_comparison.py third_party/cupy/logic_tests/test_truth.py third_party/cupy/manipulation_tests/test_basic.py diff --git a/dpnp/backend/extensions/lapack/CMakeLists.txt b/dpnp/backend/extensions/lapack/CMakeLists.txt index 7679db38d6a7..d224c623c8cb 100644 --- a/dpnp/backend/extensions/lapack/CMakeLists.txt +++ b/dpnp/backend/extensions/lapack/CMakeLists.txt @@ -27,6 +27,7 @@ set(python_module_name _lapack_impl) set(_module_src ${CMAKE_CURRENT_SOURCE_DIR}/lapack_py.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/gesv.cpp ${CMAKE_CURRENT_SOURCE_DIR}/heevd.cpp ${CMAKE_CURRENT_SOURCE_DIR}/syevd.cpp ) diff --git a/dpnp/backend/extensions/lapack/common_helpers.hpp b/dpnp/backend/extensions/lapack/common_helpers.hpp new file mode 100644 index 000000000000..2f3815320cad --- /dev/null +++ b/dpnp/backend/extensions/lapack/common_helpers.hpp @@ -0,0 +1,55 @@ +//***************************************************************************** +// Copyright (c) 2023, 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. +// +// 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 + +namespace dpnp +{ +namespace backend +{ +namespace ext +{ +namespace lapack +{ +namespace helper +{ +template +struct value_type_of +{ + using type = T; +}; + +template +struct value_type_of> +{ + using type = T; +}; +} // namespace helper +} // namespace lapack +} // namespace ext +} // namespace backend +} // namespace dpnp diff --git a/dpnp/backend/extensions/lapack/gesv.cpp b/dpnp/backend/extensions/lapack/gesv.cpp new file mode 100644 index 000000000000..72e5aa806714 --- /dev/null +++ b/dpnp/backend/extensions/lapack/gesv.cpp @@ -0,0 +1,297 @@ +//***************************************************************************** +// Copyright (c) 2023, 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. +// +// 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 + +// dpctl tensor headers +#include "utils/memory_overlap.hpp" +#include "utils/type_utils.hpp" + +#include "common_helpers.hpp" +#include "gesv.hpp" +#include "linalg_exceptions.hpp" +#include "types_matrix.hpp" + +#include "dpnp_utils.hpp" + +namespace dpnp +{ +namespace backend +{ +namespace ext +{ +namespace lapack +{ +namespace mkl_lapack = oneapi::mkl::lapack; +namespace py = pybind11; +namespace type_utils = dpctl::tensor::type_utils; + +typedef sycl::event (*gesv_impl_fn_ptr_t)(sycl::queue, + const std::int64_t, + const std::int64_t, + char *, + std::int64_t, + char *, + std::int64_t, + std::vector &, + const std::vector &); + +static gesv_impl_fn_ptr_t gesv_dispatch_vector[dpctl_td_ns::num_types]; + +template +static sycl::event gesv_impl(sycl::queue exec_q, + const std::int64_t n, + const std::int64_t nrhs, + char *in_a, + std::int64_t lda, + char *in_b, + 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::gesv_scratchpad_size(exec_q, n, nrhs, lda, ldb); + T *scratchpad = nullptr; + + std::int64_t *ipiv = nullptr; + + std::stringstream error_msg; + std::int64_t info = 0; + bool sycl_exception_caught = false; + + sycl::event gesv_event; + try { + scratchpad = sycl::malloc_device(scratchpad_size, exec_q); + ipiv = sycl::malloc_device(n, exec_q); + + gesv_event = mkl_lapack::gesv( + exec_q, + n, // The order of the matrix A (0 ≤ n). + nrhs, // The number of right-hand sides B (0 ≤ nrhs). + a, // Pointer to the square coefficient matrix A (n x n). + lda, // The leading dimension of a, must be at least max(1, n). + ipiv, // The pivot indices that define the permutation matrix P; + // row i of the matrix was interchanged with row ipiv(i), + // must be at least max(1, n). + b, // Pointer to the right hand side matrix B (n x nrhs). + ldb, // The leading dimension of b, must be at least max(1, n). + scratchpad, // Pointer to scratchpad memory to be used by MKL + // routine for storing intermediate results. + scratchpad_size, depends); + } catch (mkl_lapack::exception const &e) { + info = e.info(); + + if (info < 0) { + error_msg << "Parameter number " << -info + << " had an illegal value."; + } + else if (info > 0) { + T host_U; + exec_q.memcpy(&host_U, &a[(info - 1) * lda + info - 1], sizeof(T)) + .wait(); + + using ThresholdType = typename helper::value_type_of::type; + + const auto threshold = + std::numeric_limits::epsilon() * 100; + if (std::abs(host_U) < threshold) { + sycl::free(scratchpad, exec_q); + throw LinAlgError("The input coefficient matrix is singular."); + } + else { + error_msg << "Unexpected MKL exception caught during gesv() " + "call:\nreason: " + << e.what() << "\ninfo: " << e.info(); + } + } + 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 gesv() " + "call:\nreason: " + << e.what() << "\ninfo: " << e.info(); + } + } catch (sycl::exception const &e) { + error_msg << "Unexpected SYCL exception caught during gesv() call:\n" + << e.what(); + sycl_exception_caught = true; + } + + if (info != 0 || sycl_exception_caught) // an unexpected error occurs + { + if (scratchpad != nullptr) { + sycl::free(scratchpad, exec_q); + } + if (ipiv != nullptr) { + sycl::free(ipiv, exec_q); + } + throw std::runtime_error(error_msg.str()); + } + + sycl::event clean_up_event = exec_q.submit([&](sycl::handler &cgh) { + cgh.depends_on(gesv_event); + auto ctx = exec_q.get_context(); + cgh.host_task([ctx, scratchpad, ipiv]() { + sycl::free(scratchpad, ctx); + sycl::free(ipiv, ctx); + }); + }); + host_task_events.push_back(clean_up_event); + + return gesv_event; +} + +std::pair + gesv(sycl::queue exec_q, + dpctl::tensor::usm_ndarray coeff_matrix, + dpctl::tensor::usm_ndarray dependent_vals, + const std::vector &depends) +{ + const int coeff_matrix_nd = coeff_matrix.get_ndim(); + const int dependent_vals_nd = dependent_vals.get_ndim(); + + if (coeff_matrix_nd != 2) { + throw py::value_error("The coefficient matrix has ndim=" + + std::to_string(coeff_matrix_nd) + + ", but a 2-dimensional array is expected."); + } + + if (dependent_vals_nd > 2) { + throw py::value_error( + "The dependent values array has ndim=" + + std::to_string(dependent_vals_nd) + + ", but a 1-dimensional or a 2-dimensional array is expected."); + } + + const py::ssize_t *coeff_matrix_shape = coeff_matrix.get_shape_raw(); + const py::ssize_t *dependent_vals_shape = dependent_vals.get_shape_raw(); + + if (coeff_matrix_shape[0] != coeff_matrix_shape[1]) { + throw py::value_error("The coefficient matrix must be square," + " but got a shape of (" + + std::to_string(coeff_matrix_shape[0]) + ", " + + std::to_string(coeff_matrix_shape[1]) + ")."); + } + + // check compatibility of execution queue and allocation queue + if (!dpctl::utils::queues_are_compatible(exec_q, + {coeff_matrix, dependent_vals})) + { + throw py::value_error( + "Execution queue is not compatible with allocation queues"); + } + + auto const &overlap = dpctl::tensor::overlap::MemoryOverlap(); + if (overlap(coeff_matrix, dependent_vals)) { + throw py::value_error( + "The arrays of coefficients and dependent variables " + "are overlapping segments of memory"); + } + + bool is_coeff_matrix_f_contig = coeff_matrix.is_f_contiguous(); + if (!is_coeff_matrix_f_contig) { + throw py::value_error("The coefficient matrix " + "must be F-contiguous"); + } + + bool is_dependent_vals_f_contig = dependent_vals.is_f_contiguous(); + if (!is_dependent_vals_f_contig) { + throw py::value_error("The array of dependent variables " + "must be F-contiguous"); + } + + auto array_types = dpctl_td_ns::usm_ndarray_types(); + int coeff_matrix_type_id = + array_types.typenum_to_lookup_id(coeff_matrix.get_typenum()); + int dependent_vals_type_id = + array_types.typenum_to_lookup_id(dependent_vals.get_typenum()); + + if (coeff_matrix_type_id != dependent_vals_type_id) { + throw py::value_error("The types of the coefficient matrix and " + "dependent variables are mismatched"); + } + + gesv_impl_fn_ptr_t gesv_fn = gesv_dispatch_vector[coeff_matrix_type_id]; + if (gesv_fn == nullptr) { + throw py::value_error( + "No gesv implementation defined for the provided type " + "of the coefficient matrix."); + } + + char *coeff_matrix_data = coeff_matrix.get_data(); + char *dependent_vals_data = dependent_vals.get_data(); + + const std::int64_t n = coeff_matrix_shape[0]; + const std::int64_t m = dependent_vals_shape[0]; + const std::int64_t nrhs = + (dependent_vals_nd > 1) ? dependent_vals_shape[1] : 1; + + const std::int64_t lda = std::max(1UL, n); + const std::int64_t ldb = std::max(1UL, m); + + std::vector host_task_events; + sycl::event gesv_ev = + gesv_fn(exec_q, n, nrhs, coeff_matrix_data, lda, dependent_vals_data, + ldb, host_task_events, depends); + + sycl::event args_ev = dpctl::utils::keep_args_alive( + exec_q, {coeff_matrix, dependent_vals}, host_task_events); + + return std::make_pair(args_ev, gesv_ev); +} + +template +struct GesvContigFactory +{ + fnT get() + { + if constexpr (types::GesvTypePairSupportFactory::is_defined) { + return gesv_impl; + } + else { + return nullptr; + } + } +}; + +void init_gesv_dispatch_vector(void) +{ + dpctl_td_ns::DispatchVectorBuilder + contig; + contig.populate_dispatch_vector(gesv_dispatch_vector); +} +} // namespace lapack +} // namespace ext +} // namespace backend +} // namespace dpnp diff --git a/dpnp/backend/extensions/lapack/gesv.hpp b/dpnp/backend/extensions/lapack/gesv.hpp new file mode 100644 index 000000000000..24ac0d2e5be1 --- /dev/null +++ b/dpnp/backend/extensions/lapack/gesv.hpp @@ -0,0 +1,51 @@ +//***************************************************************************** +// Copyright (c) 2023, 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. +// +// 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 + +namespace dpnp +{ +namespace backend +{ +namespace ext +{ +namespace lapack +{ +extern std::pair + gesv(sycl::queue exec_q, + dpctl::tensor::usm_ndarray coeff_matrix, + dpctl::tensor::usm_ndarray dependent_vals, + const std::vector &depends); + +extern void init_gesv_dispatch_vector(void); +} // namespace lapack +} // namespace ext +} // namespace backend +} // namespace dpnp diff --git a/dpnp/backend/extensions/lapack/lapack_py.cpp b/dpnp/backend/extensions/lapack/lapack_py.cpp index 97b67d59e24e..c0765be7509d 100644 --- a/dpnp/backend/extensions/lapack/lapack_py.cpp +++ b/dpnp/backend/extensions/lapack/lapack_py.cpp @@ -30,7 +30,9 @@ #include #include +#include "gesv.hpp" #include "heevd.hpp" +#include "linalg_exceptions.hpp" #include "syevd.hpp" namespace lapack_ext = dpnp::backend::ext::lapack; @@ -39,6 +41,7 @@ namespace py = pybind11; // populate dispatch vectors void init_dispatch_vectors(void) { + lapack_ext::init_gesv_dispatch_vector(); lapack_ext::init_syevd_dispatch_vector(); } @@ -50,9 +53,21 @@ void init_dispatch_tables(void) PYBIND11_MODULE(_lapack_impl, m) { + // Register a custom LinAlgError exception in the dpnp.linalg submodule + py::module_ linalg_module = py::module_::import("dpnp.linalg"); + py::register_exception( + linalg_module, "LinAlgError", PyExc_ValueError); + init_dispatch_vectors(); init_dispatch_tables(); + m.def("_gesv", &lapack_ext::gesv, + "Call `gesv` from OneMKL LAPACK library to return " + "the solution of a system of linear equations with " + "a square coefficient matrix A and multiple dependent variables", + py::arg("sycl_queue"), py::arg("coeff_matrix"), + py::arg("dependent_vals"), py::arg("depends") = py::list()); + m.def("_heevd", &lapack_ext::heevd, "Call `heevd` from OneMKL LAPACK library to return " "the eigenvalues and eigenvectors of a complex Hermitian matrix", diff --git a/dpnp/backend/extensions/lapack/linalg_exceptions.hpp b/dpnp/backend/extensions/lapack/linalg_exceptions.hpp new file mode 100644 index 000000000000..083be22429c0 --- /dev/null +++ b/dpnp/backend/extensions/lapack/linalg_exceptions.hpp @@ -0,0 +1,54 @@ +//***************************************************************************** +// Copyright (c) 2023, 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. +// +// 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 + +namespace dpnp +{ +namespace backend +{ +namespace ext +{ +namespace lapack +{ +class LinAlgError : public std::exception +{ +public: + explicit LinAlgError(const char *message) : msg_(message) {} + + const char *what() const noexcept override + { + return msg_.c_str(); + } + +private: + std::string msg_; +}; +} // namespace lapack +} // namespace ext +} // namespace backend +} // namespace dpnp diff --git a/dpnp/backend/extensions/lapack/types_matrix.hpp b/dpnp/backend/extensions/lapack/types_matrix.hpp index 3cab18d3c63d..60521cb75a3c 100644 --- a/dpnp/backend/extensions/lapack/types_matrix.hpp +++ b/dpnp/backend/extensions/lapack/types_matrix.hpp @@ -43,6 +43,32 @@ namespace lapack { namespace types { +/** + * @brief A factory to define pairs of supported types for which + * MKL LAPACK library provides support in oneapi::mkl::lapack::gesv + * function. + * + * @tparam T Type of array containing the coefficient matrix A and + * the array of multiple dependent variables. Upon execution, the array of + * multiple dependent variables will be overwritten with the solution. + */ +template +struct GesvTypePairSupportFactory +{ + static constexpr bool is_defined = std::disjunction< + dpctl_td_ns::TypePairDefinedEntry, + dpctl_td_ns::TypePairDefinedEntry, + dpctl_td_ns::TypePairDefinedEntry, + T, + std::complex>, + dpctl_td_ns::TypePairDefinedEntry, + T, + std::complex>, + // fall-through + dpctl_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::heevd diff --git a/dpnp/linalg/dpnp_iface_linalg.py b/dpnp/linalg/dpnp_iface_linalg.py index c7437b30da60..c9c9c855728a 100644 --- a/dpnp/linalg/dpnp_iface_linalg.py +++ b/dpnp/linalg/dpnp_iface_linalg.py @@ -47,7 +47,12 @@ from dpnp.dpnp_utils import * from dpnp.linalg.dpnp_algo_linalg import * -from .dpnp_utils_linalg import dpnp_eigh +from .dpnp_utils_linalg import ( + check_stacked_2d, + check_stacked_square, + dpnp_eigh, + dpnp_solve, +) __all__ = [ "cholesky", @@ -62,6 +67,7 @@ "multi_dot", "norm", "qr", + "solve", "svd", ] @@ -499,6 +505,59 @@ def qr(x1, mode="reduced"): return call_origin(numpy.linalg.qr, x1, mode) +def solve(a, b): + """ + Solve a linear matrix equation, or system of linear scalar equations. + + For full documentation refer to :obj:`numpy.linalg.solve`. + + Returns + ------- + out : {(…, M,), (…, M, K)} dpnp.ndarray + Solution to the system ax = b. Returned shape is identical to b. + + Limitations + ----------- + Parameters `a` and `b` are supported as either :class:`dpnp.ndarray` + or :class:`dpctl.tensor.usm_ndarray`. + Input array data types are limited by supported DPNP :ref:`Data types`. + + See Also + -------- + :obj:`dpnp.dot` : Returns the dot product of two arrays. + + Examples + -------- + >>> import dpnp as dp + >>> a = dp.array([[1, 2], [3, 5]]) + >>> b = dp.array([1, 2]) + >>> x = dp.linalg.solve(a, b) + >>> x + array([-1., 1.]) + + Check that the solution is correct: + + >>> dp.allclose(dp.dot(a, x), b) + array([ True]) + + """ + + dpnp.check_supported_arrays_type(a, b) + check_stacked_2d(a) + check_stacked_square(a) + + if not ( + (a.ndim == b.ndim or a.ndim == b.ndim + 1) + and a.shape[:-1] == b.shape[: a.ndim - 1] + ): + raise dpnp.linalg.LinAlgError( + "a must have (..., M, M) shape and b must have (..., M) " + "or (..., M, K)" + ) + + return dpnp_solve(a, b) + + def svd(x1, full_matrices=True, compute_uv=True, hermitian=False): """ Singular Value Decomposition. diff --git a/dpnp/linalg/dpnp_utils_linalg.py b/dpnp/linalg/dpnp_utils_linalg.py index 54c01c20248e..a3a2802072c9 100644 --- a/dpnp/linalg/dpnp_utils_linalg.py +++ b/dpnp/linalg/dpnp_utils_linalg.py @@ -28,16 +28,152 @@ import dpctl.tensor._tensor_impl as ti +from numpy import issubdtype import dpnp import dpnp.backend.extensions.lapack._lapack_impl as li +from dpnp.dpnp_utils import get_usm_allocations -__all__ = ["dpnp_eigh"] +__all__ = [ + "check_stacked_2d", + "check_stacked_square", + "dpnp_eigh", + "dpnp_solve", +] _jobz = {"N": 0, "V": 1} _upper_lower = {"U": 0, "L": 1} +def check_stacked_2d(*arrays): + """ + Return ``True`` if each array in `arrays` has at least two dimensions. + + If any array is less than two-dimensional, `dpnp.linalg.LinAlgError` will be raised. + + Parameters + ---------- + arrays : {dpnp_array, usm_ndarray} + A sequence of input arrays to check for dimensionality. + + Returns + ------- + out : bool + ``True`` if each array in `arrays` is at least two-dimensional. + + Raises + ------ + dpnp.linalg.LinAlgError + If any array in `arrays` is less than two-dimensional. + + """ + + for a in arrays: + if a.ndim < 2: + raise dpnp.linalg.LinAlgError( + f"{a.ndim}-dimensional array given. The input " + "array must be at least two-dimensional" + ) + + +def check_stacked_square(*arrays): + """ + Return ``True`` if each array in `arrays` is a square matrix. + + If any array does not form a square matrix, `dpnp.linalg.LinAlgError` will be raised. + + Precondition: `arrays` are at least 2d. The caller should assert it + beforehand. For example, + + >>> def solve(a): + ... check_stacked_2d(a) + ... check_stacked_square(a) + ... ... + + Parameters + ---------- + arrays : {dpnp_array, usm_ndarray} + A sequence of input arrays to check for square matrix shape. + + Returns + ------- + out : bool + ``True`` if each array in `arrays` forms a square matrix. + + Raises + ------ + dpnp.linalg.LinAlgError + If any array in `arrays` does not form a square matrix. + + """ + + for a in arrays: + m, n = a.shape[-2:] + if m != n: + raise dpnp.linalg.LinAlgError( + "Last 2 dimensions of the input array must be square" + ) + + +def _common_type(*arrays): + """ + _common_type(*arrays) + + Common type for linear algebra operations. + + This function determines the common data type for linalg operations. + It's designed to be similar in logic to `numpy.linalg.linalg._commonType`. + + Key differences from `numpy.common_type`: + - It accepts ``bool_`` arrays. + - The default floating-point data type is determined by the capabilities of the device + on which `arrays` are created, as indicated by `dpnp.default_float_type()`. + + Args: + *arrays (dpnp.ndarray): Input arrays. + + Returns: + dtype_common (dtype): The common data type for linalg operations. + + This returned value is applicable both as the precision to be used + in linalg calls and as the dtype of (possibly complex) output(s). + + """ + + dtypes = [arr.dtype for arr in arrays] + + default = dpnp.default_float_type(device=arrays[0].device) + dtype_common = _common_inexact_type(default, *dtypes) + + return dtype_common + + +def _common_inexact_type(default_dtype, *dtypes): + """ + _common_inexact_type(default_dtype, *dtypes) + + Determines the common 'inexact' data type for linear algebra operations. + + This function selects an 'inexact' data type appropriate for the device's capabilities. + It defaults to `default_dtype` when provided types are not 'inexact'. + + Args: + default_dtype: The default data type. This is determined by the capabilities of + the device and is used when none of the provided types are 'inexact'. + *dtypes: A variable number of data types to be evaluated to find + the common 'inexact' type. + + Returns: + dpnp.result_type (dtype) : The resultant 'inexact' data type for linalg operations, + ensuring computational compatibility. + + """ + inexact_dtypes = [ + dt if issubdtype(dt, dpnp.inexact) else default_dtype for dt in dtypes + ] + return dpnp.result_type(*inexact_dtypes) + + def dpnp_eigh(a, UPLO): """ dpnp_eigh(a, UPLO) @@ -164,3 +300,135 @@ def dpnp_eigh(a, UPLO): ht_copy_ev.wait() return w, out_v + + +def dpnp_solve(a, b): + """ + dpnp_solve(a, b) + + Return the solution to the system of linear equations with + a square coefficient matrix `a` and multiple dependent variables + array `b`. + + """ + + a_usm_arr = dpnp.get_usm_ndarray(a) + b_usm_arr = dpnp.get_usm_ndarray(b) + + b_order = "C" if b.flags.c_contiguous else "F" + a_shape = a.shape + b_shape = b.shape + + res_usm_type, exec_q = get_usm_allocations([a, b]) + + res_type = _common_type(a, b) + if b.size == 0: + return dpnp.empty_like(b, dtype=res_type, usm_type=res_usm_type) + + if a.ndim > 2: + reshape = False + orig_shape_b = b_shape + if a.ndim > 3: + # get 3d input arrays by reshape + if a.ndim == b.ndim: + b = b.reshape(-1, b_shape[-2], b_shape[-1]) + else: + b = b.reshape(-1, b_shape[-1]) + + a = a.reshape(-1, a_shape[-2], a_shape[-1]) + + a_usm_arr = dpnp.get_usm_ndarray(a) + b_usm_arr = dpnp.get_usm_ndarray(b) + reshape = True + + batch_size = a.shape[0] + + coeff_vecs = [None] * batch_size + val_vecs = [None] * batch_size + a_ht_copy_ev = [None] * batch_size + b_ht_copy_ev = [None] * batch_size + ht_lapack_ev = [None] * batch_size + + for i in range(batch_size): + # oneMKL LAPACK assumes fortran-like array as input, so + # allocate a memory with 'F' order for dpnp array of coefficient matrix + # and multiple dependent variables array + coeff_vecs[i] = dpnp.empty_like( + a[i], order="F", dtype=res_type, usm_type=res_usm_type + ) + val_vecs[i] = dpnp.empty_like( + b[i], order="F", dtype=res_type, usm_type=res_usm_type + ) + + # use DPCTL tensor function to fill the coefficient matrix array + # and the array of multiple dependent variables with content + # from the input arrays + a_ht_copy_ev[i], a_copy_ev = ti._copy_usm_ndarray_into_usm_ndarray( + src=a_usm_arr[i], + dst=coeff_vecs[i].get_array(), + sycl_queue=a.sycl_queue, + ) + b_ht_copy_ev[i], b_copy_ev = ti._copy_usm_ndarray_into_usm_ndarray( + src=b_usm_arr[i], + dst=val_vecs[i].get_array(), + sycl_queue=b.sycl_queue, + ) + + # Call the LAPACK extension function _gesv to solve the system of linear + # equations using a portion of the coefficient square matrix and a + # corresponding portion of the dependent variables array. + ht_lapack_ev[i], _ = li._gesv( + exec_q, + coeff_vecs[i].get_array(), + val_vecs[i].get_array(), + depends=[a_copy_ev, b_copy_ev], + ) + + for i in range(batch_size): + ht_lapack_ev[i].wait() + b_ht_copy_ev[i].wait() + a_ht_copy_ev[i].wait() + + # combine the list of solutions into a single array + out_v = dpnp.array( + val_vecs, order=b_order, dtype=res_type, usm_type=res_usm_type + ) + if reshape: + # shape of the out_v must be equal to the shape of the array of + # dependent variables + out_v = out_v.reshape(orig_shape_b) + return out_v + else: + # oneMKL LAPACK gesv overwrites `a` and `b` and assumes fortran-like array as input. + # Allocate 'F' order memory for dpnp arrays to comply with these requirements. + a_f = dpnp.empty_like( + a, order="F", dtype=res_type, usm_type=res_usm_type + ) + + # use DPCTL tensor function to fill the coefficient matrix array + # with content from the input array `a` + a_ht_copy_ev, a_copy_ev = ti._copy_usm_ndarray_into_usm_ndarray( + src=a_usm_arr, dst=a_f.get_array(), sycl_queue=a.sycl_queue + ) + + b_f = dpnp.empty_like( + b, order="F", dtype=res_type, usm_type=res_usm_type + ) + + # use DPCTL tensor function to fill the array of multiple dependent variables + # with content from the input array `b` + b_ht_copy_ev, b_copy_ev = ti._copy_usm_ndarray_into_usm_ndarray( + src=b_usm_arr, dst=b_f.get_array(), sycl_queue=b.sycl_queue + ) + + # Call the LAPACK extension function _gesv to solve the system of linear + # equations with the coefficient square matrix and the dependent variables array. + ht_lapack_ev, _ = li._gesv( + exec_q, a_f.get_array(), b_f.get_array(), [a_copy_ev, b_copy_ev] + ) + + ht_lapack_ev.wait() + b_ht_copy_ev.wait() + a_ht_copy_ev.wait() + + return b_f diff --git a/tests/test_linalg.py b/tests/test_linalg.py index 7ce774d0abd6..6c8c2a5f04e3 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -1,11 +1,16 @@ import dpctl import numpy import pytest -from numpy.testing import assert_allclose, assert_array_equal +from numpy.testing import assert_allclose, assert_array_equal, assert_raises import dpnp as inp -from .helper import get_all_dtypes, get_complex_dtypes, has_support_aspect64 +from .helper import ( + assert_dtype_allclose, + get_all_dtypes, + has_support_aspect64, + is_cpu_device, +) def vvsort(val, vec, size, xp): @@ -508,3 +513,108 @@ def test_svd(type, shape): assert_allclose( inp.asnumpy(dpnp_vt)[i, :], np_vt[i, :], rtol=tol, atol=tol ) + + +class TestSolve: + @pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True)) + def test_solve(self, dtype): + a_np = numpy.array([[1, 0.5], [0.5, 1]], dtype=dtype) + a_dp = inp.array(a_np) + + expected = numpy.linalg.solve(a_np, a_np) + result = inp.linalg.solve(a_dp, a_dp) + + assert_allclose(expected, result, rtol=1e-06) + + @pytest.mark.parametrize("a_dtype", get_all_dtypes(no_bool=True)) + @pytest.mark.parametrize("b_dtype", get_all_dtypes(no_bool=True)) + def test_solve_diff_type(self, a_dtype, b_dtype): + a_np = numpy.array([[1, 2], [3, -5]], dtype=a_dtype) + b_np = numpy.array([4, 1], dtype=b_dtype) + + a_dp = inp.array(a_np) + b_dp = inp.array(b_np) + + expected = numpy.linalg.solve(a_np, b_np) + result = inp.linalg.solve(a_dp, b_dp) + + assert_dtype_allclose(result, expected) + + def test_solve_strides(self): + a_np = numpy.array( + [ + [2, 3, 1, 4, 5], + [5, 6, 7, 8, 9], + [9, 7, 7, 2, 3], + [1, 4, 5, 1, 8], + [8, 9, 8, 5, 3], + ] + ) + b_np = numpy.array([5, 8, 9, 2, 1]) + + a_dp = inp.array(a_np) + b_dp = inp.array(b_np) + + # positive strides + expected = numpy.linalg.solve(a_np[::2, ::2], b_np[::2]) + result = inp.linalg.solve(a_dp[::2, ::2], b_dp[::2]) + assert_allclose(expected, result, rtol=1e-05) + + # negative strides + expected = numpy.linalg.solve(a_np[::-2, ::-2], b_np[::-2]) + result = inp.linalg.solve(a_dp[::-2, ::-2], b_dp[::-2]) + assert_allclose(expected, result, rtol=1e-05) + + # TODO: remove skipif when MKLD-16626 is resolved + @pytest.mark.skipif(is_cpu_device(), reason="MKLD-16626") + @pytest.mark.parametrize( + "matrix, vector", + [ + ([[1, 2], [2, 4]], [1, 2]), + ([[0, 0], [0, 0]], [0, 0]), + ([[1, 1], [1, 1]], [2, 2]), + ([[2, 4], [1, 2]], [3, 1.5]), + ([[1, 2], [0, 0]], [3, 0]), + ([[1, 0], [2, 0]], [3, 4]), + ], + ids=[ + "Linearly dependent rows", + "Zero matrix", + "Identical rows", + "Linearly dependent columns", + "Zero row", + "Zero column", + ], + ) + def test_solve_singular_matrix(self, matrix, vector): + a_np = numpy.array(matrix, dtype="float32") + b_np = numpy.array(vector, dtype="float32") + + a_dp = inp.array(a_np) + b_dp = inp.array(b_np) + + assert_raises(numpy.linalg.LinAlgError, numpy.linalg.solve, a_np, b_np) + assert_raises(inp.linalg.LinAlgError, inp.linalg.solve, a_dp, b_dp) + + def test_solve_errors(self): + a_dp = inp.array([[1, 0.5], [0.5, 1]], dtype="float32") + b_dp = inp.array(a_dp, dtype="float32") + + # diffetent queue + a_queue = dpctl.SyclQueue() + b_queue = dpctl.SyclQueue() + a_dp_q = inp.array(a_dp, sycl_queue=a_queue) + b_dp_q = inp.array(b_dp, sycl_queue=b_queue) + assert_raises(ValueError, inp.linalg.solve, a_dp_q, b_dp_q) + + # unsupported type + a_np = inp.asnumpy(a_dp) + b_np = inp.asnumpy(b_dp) + assert_raises(TypeError, inp.linalg.solve, a_np, b_dp) + assert_raises(TypeError, inp.linalg.solve, a_dp, b_np) + + # a.ndim < 2 + a_dp_ndim_1 = a_dp.flatten() + assert_raises( + inp.linalg.LinAlgError, inp.linalg.solve, a_dp_ndim_1, b_dp + ) diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index 3c658c14fe52..19104d7bd434 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -1345,3 +1345,27 @@ def test_take(func, device): expected_queue = dpnp_data.get_array().sycl_queue result_queue = result.get_array().sycl_queue assert_sycl_queue_equal(result_queue, expected_queue) + + +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +def test_solve(device): + x = [[1.0, 2.0], [3.0, 5.0]] + y = [1.0, 2.0] + + numpy_x = numpy.array(x) + numpy_y = numpy.array(y) + dpnp_x = dpnp.array(x, device=device) + dpnp_y = dpnp.array(y, device=device) + + result = dpnp.linalg.solve(dpnp_x, dpnp_y) + expected = numpy.linalg.solve(numpy_x, numpy_y) + assert_dtype_allclose(result, expected) + + result_queue = result.sycl_queue + + assert_sycl_queue_equal(result_queue, dpnp_x.sycl_queue) + assert_sycl_queue_equal(result_queue, dpnp_y.sycl_queue) diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 4982ed424140..fa2931cc5054 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -486,3 +486,41 @@ def test_take(func, usm_type_x, usm_type_ind): assert x.usm_type == usm_type_x assert ind.usm_type == usm_type_ind assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_ind]) + + +@pytest.mark.parametrize( + "usm_type_matrix", list_of_usm_types, ids=list_of_usm_types +) +@pytest.mark.parametrize( + "usm_type_vector", list_of_usm_types, ids=list_of_usm_types +) +@pytest.mark.parametrize( + "matrix, vector", + [ + ([[1, 2], [3, 5]], dp.empty((2, 0))), + ([[1, 2], [3, 5]], [1, 2]), + ( + [ + [[1, 1, 1], [0, 2, 5], [2, 5, -1]], + [[3, -1, 1], [1, 2, 3], [2, 3, 1]], + [[1, 4, 1], [1, 2, -2], [4, 1, 2]], + ], + [[6, -4, 27], [9, -6, 15], [15, 1, 11]], + ), + ], + ids=[ + "2D_Matrix_Empty_Vector", + "2D_Matrix_1D_Vector", + "3D_Matrix_and_Vectors", + ], +) +def test_solve(matrix, vector, usm_type_matrix, usm_type_vector): + x = dp.array(matrix, usm_type=usm_type_matrix) + y = dp.array(vector, usm_type=usm_type_vector) + z = dp.linalg.solve(x, y) + + assert x.usm_type == usm_type_matrix + assert y.usm_type == usm_type_vector + assert z.usm_type == du.get_coerced_usm_type( + [usm_type_matrix, usm_type_vector] + ) diff --git a/tests/third_party/cupy/linalg_tests/test_solve.py b/tests/third_party/cupy/linalg_tests/test_solve.py new file mode 100644 index 000000000000..6194cf6b8ac4 --- /dev/null +++ b/tests/third_party/cupy/linalg_tests/test_solve.py @@ -0,0 +1,90 @@ +import unittest + +import numpy +import pytest + +import dpnp as cupy +from tests.helper import has_support_aspect64 +from tests.third_party.cupy import testing + + +@testing.parameterize( + *testing.product( + { + "order": ["C", "F"], + } + ) +) +class TestSolve(unittest.TestCase): + # TODO: add get_batched_gesv_limit + # def setUp(self): + # if self.batched_gesv_limit is not None: + # self.old_limit = get_batched_gesv_limit() + # set_batched_gesv_limit(self.batched_gesv_limit) + + # def tearDown(self): + # if self.batched_gesv_limit is not None: + # set_batched_gesv_limit(self.old_limit) + + @testing.for_dtypes("ifdFD") + @testing.numpy_cupy_allclose( + atol=1e-3, contiguous_check=False, type_check=has_support_aspect64() + ) + def check_x(self, a_shape, b_shape, xp, dtype): + a = testing.shaped_random(a_shape, xp, dtype=dtype, seed=0, scale=20) + b = testing.shaped_random(b_shape, xp, dtype=dtype, seed=1) + a = a.copy(order=self.order) + b = b.copy(order=self.order) + a_copy = a.copy() + b_copy = b.copy() + result = xp.linalg.solve(a, b) + numpy.testing.assert_array_equal(a_copy, a) + numpy.testing.assert_array_equal(b_copy, b) + return result + + def test_solve(self): + self.check_x((4, 4), (4,)) + self.check_x((5, 5), (5, 2)) + self.check_x((2, 4, 4), (2, 4)) + self.check_x((2, 5, 5), (2, 5, 2)) + self.check_x((2, 3, 2, 2), (2, 3, 2)) + self.check_x((2, 3, 3, 3), (2, 3, 3, 2)) + self.check_x((0, 0), (0,)) + self.check_x((0, 0), (0, 2)) + self.check_x((0, 2, 2), (0, 2)) + self.check_x((0, 2, 2), (0, 2, 3)) + + def check_shape(self, a_shape, b_shape, error_types): + for xp, error_type in error_types.items(): + a = xp.random.rand(*a_shape) + b = xp.random.rand(*b_shape) + with pytest.raises(error_type): + xp.linalg.solve(a, b) + + # Undefined behavior is implementation-dependent: + # Numpy with OpenBLAS returns an empty array + # while numpy with OneMKL raises LinAlgError + @pytest.mark.skip("Undefined behavior") + def test_solve_singular_empty(self, xp): + a = xp.zeros((3, 3)) # singular + b = xp.empty((3, 0)) # nrhs = 0 + # LinAlgError("Singular matrix") is not raised + return xp.linalg.solve(a, b) + + def test_invalid_shape(self): + linalg_errors = { + numpy: numpy.linalg.LinAlgError, + cupy: cupy.linalg.LinAlgError, + } + value_errors = { + numpy: ValueError, + cupy: ValueError, + } + + self.check_shape((2, 3), (4,), linalg_errors) + self.check_shape((3, 3), (2,), value_errors) + self.check_shape((3, 3), (2, 2), value_errors) + self.check_shape((3, 3, 4), (3,), linalg_errors) + self.check_shape((2, 3, 3), (3,), value_errors) + self.check_shape((3, 3), (0,), value_errors) + self.check_shape((0, 3, 4), (3,), linalg_errors)