diff --git a/tests/linear_algebra/cases.hpp b/tests/linear_algebra/cases.hpp index 70c756e..4934f3b 100644 --- a/tests/linear_algebra/cases.hpp +++ b/tests/linear_algebra/cases.hpp @@ -157,8 +157,124 @@ int run_linear_algebra_type() { return 0; } +template +int run_cholesky_type() { + constexpr auto order{3uz}; + constexpr auto stride{5uz}; + constexpr auto size{order * stride}; + constexpr auto padding{T{19}}; + constexpr auto tolerance{T{100} * std::numeric_limits::epsilon()}; + + constexpr auto original = std::array{ + T{4}, T{2}, T{2}, + T{2}, T{5}, T{1}, + T{2}, T{1}, T{6} + }; + + auto matrix_host = std::array{}; + matrix_host.fill(padding); + for (auto row{0uz}; row < order; ++row) { + for (auto column{0uz}; column < order; ++column) { + matrix_host[row * stride + column] = original[row * order + column]; + } + } + + auto matrix{xpu::buffer{size}}; + xpu::copy_n(matrix.data(), matrix_host.data(), size); + + auto factorization{xpu::linalg::cholesky_factorization{order, stride}}; + if ( + factorization.order() != order || + factorization.stride() != stride + ) { + return test::fail("cholesky_factorization shape accessors returned the wrong values"); + } + + if ( + factorization.factorize(matrix.data()) != + xpu::linalg::status::success + ) { + return test::fail("cholesky_factorization reported a positive-definite matrix as not positive definite"); + } + + xpu::copy_n(matrix_host.data(), matrix.data(), size); + auto diagonal_product{T{1}}; + for (auto i{0uz}; i < order; ++i) { + diagonal_product *= matrix_host[i * stride + i]; + for (auto column{order}; column < stride; ++column) { + if (matrix_host[i * stride + column] != padding) { + return test::fail("Cholesky factorize modified row padding"); + } + } + } + + if (!test::near(diagonal_product * diagonal_product, T{80}, tolerance)) { + return test::fail("Cholesky diagonal has the wrong determinant"); + } + + constexpr auto rhs_host = std::array{T{14}, T{15}, T{22}}; + constexpr auto expected_solution = std::array{T{1}, T{2}, T{3}}; + auto solution_host = std::array{}; + auto rhs{xpu::buffer{order}}; + auto solution{xpu::buffer{order}}; + xpu::copy_n(rhs.data(), rhs_host.data(), order); + + factorization.solve(matrix.data(), rhs.data(), solution.data()); + xpu::copy_n(solution_host.data(), solution.data(), order); + + for (auto i{0uz}; i < order; ++i) { + if (!test::near(solution_host[i], expected_solution[i], tolerance)) { + return test::fail("Cholesky solve produced the wrong solution"); + } + } + + matrix_host.fill(padding); + constexpr auto not_positive_definite = std::array{ + T{1}, T{0}, T{0}, + T{0}, T{1}, T{0}, + T{0}, T{0}, T{-1} + }; + for (auto row{0uz}; row < order; ++row) { + for (auto column{0uz}; column < order; ++column) { + matrix_host[row * stride + column] = + not_positive_definite[row * order + column]; + } + } + xpu::copy_n(matrix.data(), matrix_host.data(), size); + + if ( + factorization.factorize(matrix.data()) != + xpu::linalg::status::not_pd + ) { + return test::fail("cholesky_factorization did not report a non-positive-definite matrix"); + } + + return 0; +} + inline auto run_linear_algebra_checked_cases() -> int { #if defined(__unix__) + if (const auto failure{test::check_abort([] { + auto factorization{xpu::linalg::cholesky_factorization{std::numeric_limits::max(), 1uz}}; + })}; failure) { + + return failure; + } + + if (const auto failure{test::check_abort([] { + auto factorization{xpu::linalg::cholesky_factorization{1uz, std::numeric_limits::max()}}; + })}; failure) { + + return failure; + } + + if (const auto failure{test::check_abort([] { + auto factorization{xpu::linalg::cholesky_factorization{2uz, 1uz}}; + })}; failure) { + + return failure; + } + if (const auto failure{test::check_abort([] { auto factorization{xpu::linalg::lu_factorization{std::numeric_limits::max(), 1uz}}; })}; failure) { @@ -217,5 +333,11 @@ inline int run_linear_algebra_cases() { if (const auto status{run_linear_algebra_type()}; status != 0) { return status; } - return run_linear_algebra_type(); + if (const auto status{run_linear_algebra_type()}; status != 0) { + return status; + } + if (const auto status{run_cholesky_type()}; status != 0) { + return status; + } + return run_cholesky_type(); } diff --git a/xpu/linear_algebra.hpp b/xpu/linear_algebra.hpp index 38a918a..af246b4 100644 --- a/xpu/linear_algebra.hpp +++ b/xpu/linear_algebra.hpp @@ -24,7 +24,8 @@ namespace linalg { enum class status { success, - singular + singular, + not_pd }; namespace detail { @@ -122,6 +123,38 @@ inline auto getrf_workspace_size( return workspace_size; } +template +inline auto potrf_workspace_size( + cusolverDnHandle_t handle, + std::size_t order, + std::size_t stride +) -> std::size_t { + const auto vendor_order{xpu::detail::checked_cast(order)}; + const auto vendor_stride{xpu::detail::checked_cast(stride)}; + + auto size{0}; + + if constexpr (std::same_as) { + xpu::cu_check(cusolverDnSpotrf_bufferSize( + handle, CUBLAS_FILL_MODE_UPPER, + vendor_order, + nullptr, vendor_stride, + &size + )); + } else { + xpu::cu_check(cusolverDnDpotrf_bufferSize( + handle, CUBLAS_FILL_MODE_UPPER, + vendor_order, + nullptr, vendor_stride, + &size + )); + } + + const auto workspace_size{xpu::detail::checked_cast(size)}; + + return workspace_size; +} + template inline auto cusolver_getrf( cusolverDnHandle_t handle, @@ -152,6 +185,39 @@ inline auto cusolver_getrf( } } +template +inline auto cusolver_potrf( + cusolverDnHandle_t handle, + std::size_t order, + std::size_t stride, + T* matrix, + T* workspace, + std::size_t workspace_size, + int* info +) -> void { + const auto vendor_order{xpu::detail::checked_cast(order)}; + const auto vendor_stride{xpu::detail::checked_cast(stride)}; + const auto vendor_workspace_size{xpu::detail::checked_cast(workspace_size)}; + + if constexpr (std::same_as) { + xpu::cu_check(cusolverDnSpotrf( + handle, CUBLAS_FILL_MODE_UPPER, + vendor_order, + matrix, vendor_stride, + workspace, vendor_workspace_size, + info + )); + } else { + xpu::cu_check(cusolverDnDpotrf( + handle, CUBLAS_FILL_MODE_UPPER, + vendor_order, + matrix, vendor_stride, + workspace, vendor_workspace_size, + info + )); + } +} + template inline auto cusolver_getrs( cusolverDnHandle_t handle, @@ -191,6 +257,41 @@ inline auto cusolver_getrs( } } +template +inline auto cusolver_potrs( + cusolverDnHandle_t handle, + std::size_t order, + std::size_t right_hand_sides, + const T* factor, + std::size_t factor_stride, + T* solution, + std::size_t solution_stride, + int* info +) -> void { + const auto vendor_order{xpu::detail::checked_cast(order)}; + const auto vendor_right_hand_sides{xpu::detail::checked_cast(right_hand_sides)}; + const auto vendor_factor_stride{xpu::detail::checked_cast(factor_stride)}; + const auto vendor_solution_stride{xpu::detail::checked_cast(solution_stride)}; + + if constexpr (std::same_as) { + xpu::cu_check(cusolverDnSpotrs( + handle, CUBLAS_FILL_MODE_UPPER, + vendor_order, vendor_right_hand_sides, + factor, vendor_factor_stride, + solution, vendor_solution_stride, + info + )); + } else { + xpu::cu_check(cusolverDnDpotrs( + handle, CUBLAS_FILL_MODE_UPPER, + vendor_order, vendor_right_hand_sides, + factor, vendor_factor_stride, + solution, vendor_solution_stride, + info + )); + } +} + #else template @@ -220,6 +321,30 @@ inline auto lapacke_getrf( } } +template +inline auto lapacke_potrf( + T* RESTRICT matrix, + std::size_t order, + std::size_t stride +) -> lapack_int { + const auto vendor_order{xpu::detail::checked_cast(order)}; + const auto vendor_stride{xpu::detail::checked_cast(stride)}; + + if constexpr (std::same_as) { + return LAPACKE_spotrf( + LAPACK_ROW_MAJOR, 'L', + vendor_order, + matrix, vendor_stride + ); // use lower triangular + } else { + return LAPACKE_dpotrf( + LAPACK_ROW_MAJOR, 'L', + vendor_order, + matrix, vendor_stride + ); + } +} + template inline auto lapacke_getri( T* RESTRICT inverse, @@ -277,6 +402,33 @@ inline auto lapacke_getrs( } } +template +inline auto lapacke_potrs( + const T* RESTRICT lower_upper, + T* RESTRICT solution, + std::size_t order, + std::size_t stride +) -> lapack_int { + const auto vendor_order{xpu::detail::checked_cast(order)}; + const auto vendor_stride{xpu::detail::checked_cast(stride)}; + + if constexpr (std::same_as) { + return LAPACKE_spotrs( + LAPACK_ROW_MAJOR, 'L', + vendor_order, 1, + lower_upper, vendor_stride, + solution, 1 + ); + } else { + return LAPACKE_dpotrs( + LAPACK_ROW_MAJOR, 'L', + vendor_order, 1, + lower_upper, vendor_stride, + solution, 1 + ); + } +} + #endif } // namespace xpu::linalg::detail @@ -332,6 +484,148 @@ inline auto transpose_square( #endif } +template +class cholesky_factorization { +private: + std::size_t order_; + std::size_t stride_; + T* factor_{}; + +#if defined(XPU_CUDA) + using dimension_type = int; +#else + using dimension_type = lapack_int; +#endif + +#if defined(XPU_CUDA) + cusolverDnHandle_t handle_; + xpu::buffer workspace_; + xpu::buffer info_; +#endif + + [[nodiscard]] + static auto checked_order(std::size_t order, std::size_t stride) noexcept -> std::size_t { + static_cast(xpu::detail::checked_cast(order)); + static_cast(xpu::detail::checked_cast(stride)); + + const auto invalid_stride{stride < order}; + + if (invalid_stride) { + detail::linalg_error("matrix stride is smaller than its order"); + } + + const auto matrix_size{xpu::detail::checked_mul(order, stride)}; + + static_cast(xpu::detail::checked_bytes(matrix_size)); + + return order; + } + +public: + cholesky_factorization(std::size_t order, std::size_t stride) + : order_{checked_order(order, stride)} + , stride_{stride} +#if defined(XPU_CUDA) + , handle_{detail::create_cusolver_handle()} + , workspace_{detail::potrf_workspace_size(handle_, order, stride)} + , info_{1uz} +#endif + { } + + ~cholesky_factorization() { +#if defined(XPU_CUDA) + xpu::cu_check(cusolverDnDestroy(handle_)); +#endif + } + + [[nodiscard]] + constexpr auto order() const noexcept -> std::size_t { + return order_; + } + + [[nodiscard]] + constexpr auto stride() const noexcept -> std::size_t { + return stride_; + } + + [[nodiscard]] + auto factorize(T* RESTRICT matrix) noexcept -> status { + factor_ = nullptr; + +#if defined(XPU_CUDA) + detail::cusolver_potrf( + handle_, order_, stride_, + matrix, workspace_.data(), workspace_.size(), info_.data() + ); + + auto info{0}; + xpu::copy_n(&info, info_.data(), 1uz); + if (info < 0) { + detail::linalg_error("cuSOLVER potrf received an invalid argument"); + } + if (info > 0) { return status::not_pd; } +#else + const auto info{ + detail::lapacke_potrf(matrix, order_, stride_) + }; + if (info < 0) { + detail::linalg_error("LAPACKE potrf failed"); + } + if (info > 0) { return status::not_pd; } +#endif + + factor_ = matrix; + return status::success; + } + + auto solve( + const T* RESTRICT factor, + const T* RESTRICT rhs, + T* RESTRICT solution + ) noexcept -> void { + if (!factor_) { + detail::linalg_error("you must factor before you solve"); + } + if (factor != factor_) { + detail::linalg_error( + "solve requires the most recently factorized matrix" + ); + } + if (rhs == solution) { + detail::linalg_error("rhs and solution must not alias"); + } + + xpu::copy_n(solution, rhs, order_); + +#if defined(XPU_CUDA) + detail::cusolver_potrs( + handle_, order_, 1uz, + factor, stride_, + solution, order_, + info_.data() + ); + + auto info{0}; + xpu::copy_n(&info, info_.data(), 1uz); + if (info != 0) { + detail::linalg_error("cuSOLVER potrs received an invalid argument"); + } +#else + const auto info{ + detail::lapacke_potrs(factor, solution, order_, stride_) + }; + if (info != 0) { + detail::linalg_error("LAPACKE potrs failed"); + } +#endif + } + + auto operator=(const cholesky_factorization&) -> cholesky_factorization& = delete; + cholesky_factorization(const cholesky_factorization&) = delete; + auto operator=(cholesky_factorization&&) -> cholesky_factorization& = delete; + cholesky_factorization(cholesky_factorization&&) = delete; +}; + template class lu_factorization { private: