perf(shell): route local products through MKL and reduce Hessian storage work
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
#include "math/dense_blas_internal.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
#include <stdexcept>
|
||||
@@ -26,8 +27,22 @@ ScopedSingleEquationThreads::~ScopedSingleEquationThreads() {
|
||||
if (restore_default_) static_cast<void>(mkl_set_num_threads_local(0));
|
||||
}
|
||||
|
||||
ScopedLocalDenseThreads::ScopedLocalDenseThreads() noexcept {
|
||||
const int previous = mkl_set_num_threads_local(1);
|
||||
if (previous == 0) {
|
||||
restore_default_ = true;
|
||||
} else {
|
||||
static_cast<void>(mkl_set_num_threads_local(previous));
|
||||
}
|
||||
}
|
||||
|
||||
ScopedLocalDenseThreads::~ScopedLocalDenseThreads() {
|
||||
if (restore_default_) static_cast<void>(mkl_set_num_threads_local(0));
|
||||
}
|
||||
|
||||
bool CholeskyLowerAndInverse(const Matrix& matrix, Matrix& lower,
|
||||
Matrix& inverse_lower) {
|
||||
Matrix& inverse_lower, CholeskyFailure* failure) {
|
||||
if (failure != nullptr) *failure = CholeskyFailure::kNone;
|
||||
if (matrix.Rows() != matrix.Columns())
|
||||
throw std::invalid_argument{"Cholesky requires a square matrix."};
|
||||
const auto size = ToMklSize(matrix.Rows());
|
||||
@@ -48,17 +63,30 @@ bool CholeskyLowerAndInverse(const Matrix& matrix, Matrix& lower,
|
||||
values[row * matrix.Columns() + column] = matrix(row, column);
|
||||
}
|
||||
if (LAPACKE_dpotrf(LAPACK_ROW_MAJOR, 'L', size.Value(), values.data(),
|
||||
size.Value()) != 0) return false;
|
||||
size.Value()) != 0) {
|
||||
if (failure != nullptr) *failure = CholeskyFailure::kPivot;
|
||||
return false;
|
||||
}
|
||||
for (std::size_t row = 0; row < matrix.Rows(); ++row)
|
||||
for (std::size_t column = 0; column <= row; ++column)
|
||||
for (std::size_t column = 0; column <= row; ++column) {
|
||||
candidate_lower(row, column) = values[row * matrix.Columns() + column];
|
||||
if (!std::isfinite(candidate_lower(row, column))) {
|
||||
if (failure != nullptr) *failure = CholeskyFailure::kPivot;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if (LAPACKE_dtrtri(LAPACK_ROW_MAJOR, 'L', 'N', size.Value(), values.data(),
|
||||
size.Value()) != 0) return false;
|
||||
size.Value()) != 0) {
|
||||
if (failure != nullptr) *failure = CholeskyFailure::kInverse;
|
||||
return false;
|
||||
}
|
||||
for (std::size_t row = 0; row < matrix.Rows(); ++row)
|
||||
for (std::size_t column = 0; column <= row; ++column) {
|
||||
const auto value = values[row * matrix.Columns() + column];
|
||||
if (!std::isfinite(candidate_lower(row, column)) || !std::isfinite(value))
|
||||
if (!std::isfinite(value)) {
|
||||
if (failure != nullptr) *failure = CholeskyFailure::kInverse;
|
||||
return false;
|
||||
}
|
||||
candidate_inverse(row, column) = value;
|
||||
}
|
||||
lower = std::move(candidate_lower);
|
||||
@@ -96,6 +124,41 @@ Matrix MultiplyTransposeLeft(const Matrix& left, const Matrix& right) {
|
||||
return transposed.Multiply(right);
|
||||
}
|
||||
|
||||
bool OuterProduct(const double* left, const std::size_t rows,
|
||||
const double* right, const std::size_t columns,
|
||||
double* destination) {
|
||||
return MultiplyTransposeRight(left, rows, right, columns, 1U, destination);
|
||||
}
|
||||
|
||||
bool MultiplyTransposeRight(const double* left, const std::size_t rows,
|
||||
const double* right, const std::size_t columns,
|
||||
const std::size_t inner, double* destination) {
|
||||
const auto m = ToMklSize(rows);
|
||||
const auto n = ToMklSize(columns);
|
||||
const auto k = ToMklSize(inner);
|
||||
const auto maximum = std::numeric_limits<std::size_t>::max() / sizeof(double);
|
||||
if (!m.HasValue() || !n.HasValue() || !k.HasValue() ||
|
||||
(columns != 0U && rows > maximum / columns) ||
|
||||
(inner != 0U && (rows > maximum / inner || columns > maximum / inner)))
|
||||
throw std::length_error{"Dense product dimensions exceed backend limits."};
|
||||
if (rows == 0U || columns == 0U) return true;
|
||||
if (destination == nullptr ||
|
||||
(inner != 0U && (left == nullptr || right == nullptr)))
|
||||
throw std::invalid_argument{"Dense product requires nonnull storage."};
|
||||
if (inner == 0U) {
|
||||
std::fill_n(destination, rows * columns, 0.0);
|
||||
return true;
|
||||
}
|
||||
for (std::size_t i = 0U; i < rows * inner; ++i)
|
||||
if (!std::isfinite(left[i])) return false;
|
||||
for (std::size_t i = 0U; i < columns * inner; ++i)
|
||||
if (!std::isfinite(right[i])) return false;
|
||||
cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m.Value(), n.Value(),
|
||||
k.Value(), 1.0, left, k.Value(), right, k.Value(), 0.0,
|
||||
destination, n.Value());
|
||||
return true;
|
||||
}
|
||||
|
||||
Result<MKL_INT> ToMklSize(const std::size_t size) {
|
||||
if (size > static_cast<std::size_t>((std::numeric_limits<MKL_INT>::max)())) {
|
||||
return Result<MKL_INT>::Failure(Status::Failure(
|
||||
|
||||
@@ -26,6 +26,20 @@ class ScopedSingleEquationThreads {
|
||||
bool restore_default_{false};
|
||||
};
|
||||
|
||||
/// @brief Limits element-local dense work under the caller's default policy.
|
||||
/// @note Positive caller TLS overrides are retained. Default requests are
|
||||
/// restored on success and unwind; process-wide settings are never changed.
|
||||
class ScopedLocalDenseThreads {
|
||||
public:
|
||||
ScopedLocalDenseThreads() noexcept;
|
||||
ScopedLocalDenseThreads(const ScopedLocalDenseThreads&) = delete;
|
||||
ScopedLocalDenseThreads& operator=(const ScopedLocalDenseThreads&) = delete;
|
||||
~ScopedLocalDenseThreads();
|
||||
|
||||
private:
|
||||
bool restore_default_{false};
|
||||
};
|
||||
|
||||
/// @brief Converts a dense storage length to the private MKL integer type.
|
||||
/// @return The converted length or an uncategorized overflow failure.
|
||||
Result<MKL_INT> ToMklSize(std::size_t size);
|
||||
@@ -39,13 +53,38 @@ Matrix MultiplyTransposeLeft(const Matrix& left, const Matrix& right);
|
||||
/// @throws std::length_error if the size exceeds the backend integer range.
|
||||
Vector MultiplyDiagonal(const Vector& diagonal, const Vector& rhs);
|
||||
|
||||
/// @brief Writes one row-major outer product through MKL without allocation.
|
||||
/// @note Source arrays and destination must not alias. A nonfinite input
|
||||
/// returns false before writing; finite product overflow remains for the
|
||||
/// owner's gate. A zero dimension accepts null pointers and performs no backend
|
||||
/// call.
|
||||
/// @throws std::invalid_argument for null storage of a nonempty product.
|
||||
/// @throws std::length_error for dimensions outside the backend integer range.
|
||||
bool OuterProduct(const double* left, std::size_t rows, const double* right,
|
||||
std::size_t columns, double* destination);
|
||||
|
||||
/// @brief Writes left times right-transpose to nonaliased contiguous storage.
|
||||
/// @note Inputs are row-major rows-by-inner and columns-by-inner. Nonfinite
|
||||
/// input returns false before writing. Finite output overflow is not hidden.
|
||||
/// Zero output dimensions accept null pointers; inner zero writes exact zeros.
|
||||
/// @throws std::invalid_argument for missing nonempty storage.
|
||||
/// @throws std::length_error for dimension or storage-size overflow.
|
||||
bool MultiplyTransposeRight(const double* left, std::size_t rows,
|
||||
const double* right, std::size_t columns,
|
||||
std::size_t inner, double* destination);
|
||||
|
||||
/// @brief Distinguishes positive-factor failure from inverse failure.
|
||||
enum class CholeskyFailure { kNone, kPivot, kInverse };
|
||||
|
||||
/// @brief Computes a lower Cholesky factor and its inverse with MKL LAPACK.
|
||||
/// @return False for a nonpositive pivot or nonfinite factor/inverse.
|
||||
/// @note Outputs are replaced only after both finite candidates succeed.
|
||||
/// The optional failure output is initialized even when validation throws.
|
||||
/// @throws std::invalid_argument for a nonsquare or nonfinite input.
|
||||
/// @throws std::length_error for a backend size overflow.
|
||||
bool CholeskyLowerAndInverse(const Matrix& matrix, Matrix& lower,
|
||||
Matrix& inverse_lower);
|
||||
Matrix& inverse_lower,
|
||||
CholeskyFailure* failure = nullptr);
|
||||
|
||||
/// @brief Copies contiguous double values through the private BLAS backend.
|
||||
/// @param source Readable storage for size values when size is nonzero.
|
||||
|
||||
Reference in New Issue
Block a user