diff --git a/.gitignore b/.gitignore index eb5e6dd71..cf2691220 100644 --- a/.gitignore +++ b/.gitignore @@ -66,3 +66,6 @@ benchmark/build/** test/test_runner z_ignore_riley/** +benchmark-build-audit/ +build-audit/ +install-audit/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 20687ffc7..181dd595c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -44,8 +44,6 @@ list(APPEND CMAKE_MODULE_PATH ${CMAKE_SOURCE_DIR}) include(compiler_flags) # Configure the build -enable_testing() - include(rl_build_options) include(rl_version) @@ -60,6 +58,7 @@ add_subdirectory(RandBLAS) # Compile sources add_subdirectory(RandLAPACK) if (RandLAPACK_BUILD_TESTS) + enable_testing() add_subdirectory(test) endif() diff --git a/RandLAPACK.hh b/RandLAPACK.hh index 7e88f5cf2..692aeb577 100644 --- a/RandLAPACK.hh +++ b/RandLAPACK.hh @@ -28,12 +28,19 @@ #include "RandLAPACK/comps/rl_syrf.hh" #include "RandLAPACK/comps/rl_orth.hh" #include "RandLAPACK/comps/rl_rpchol.hh" +#include "RandLAPACK/comps/rl_cholqr.hh" // Drivers #include "RandLAPACK/drivers/rl_rsvd.hh" -#include "RandLAPACK/drivers/rl_cqrrt.hh" +#include "RandLAPACK/drivers/rl_cqrrt.hh" // holds both dense CQRRT and CQRRT_linops #include "RandLAPACK/drivers/rl_cholqr_linops.hh" -#include "RandLAPACK/drivers/rl_cqrrt_linops.hh" +#include "RandLAPACK/drivers/rl_cholqr_dense.hh" +#include "RandLAPACK/drivers/rl_iter_refine_lsq.hh" +// Both of these declare themselves "Public API" in their headers but were reachable +// only by including them directly, which is part of why neither had any test coverage. +#include "RandLAPACK/drivers/rl_lsqr.hh" +#include "RandLAPACK/drivers/rl_restarted_pcg_ne.hh" +#include "RandLAPACK/drivers/rl_blendenpik.hh" #include "RandLAPACK/drivers/rl_scholqr3_linops.hh" #include "RandLAPACK/drivers/rl_cqrrpt.hh" #include "RandLAPACK/drivers/rl_bqrrp.hh" diff --git a/RandLAPACK/comps/rl_cholqr.hh b/RandLAPACK/comps/rl_cholqr.hh new file mode 100644 index 000000000..b13cb5031 --- /dev/null +++ b/RandLAPACK/comps/rl_cholqr.hh @@ -0,0 +1,747 @@ +#pragma once + +#include "rl_util.hh" +#include "rl_blaspp.hh" +#include "rl_lapackpp.hh" +#include "rl_linops.hh" +#include "rl_bqrrp.hh" +#include "rl_exceptions.hh" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace RandLAPACK { + +/// Below enum lists the methods cholqr_primitive can use to invert the +/// upper-triangular preconditioner P (forming R_pre = P^{-1}) in the +/// preconditioned path. They trade speed for stability when P is ill-conditioned: +/// +/// TRSM_IDENTITY Solve P * R_pre = I via TRSM(Side::Left). Backward-stable for the +/// subsequent product A * R_pre. O(n^3/2). Default. +/// TRTRI Direct triangular inverse via LAPACK trtri. Same cost; less stable +/// than TRSM_IDENTITY when P is ill-conditioned in practice. +/// GEQP3 R_pre via column-pivoted QR of P (geqp3), giving P^{-1} = Pi R^{-1} Q^T. +/// Stable for ill-conditioned P. O(11 n^3 / 6). +/// BQRRP Same output as GEQP3 but pivots via RandLAPACK::BQRRP (blocked + +/// sketched). Faster for large n; requires an RNG state. +enum class PCholQRPrecondMethod { + TRSM_IDENTITY, + TRTRI, + GEQP3, + BQRRP +}; + + +// Default column-block width for the blocked Gram computation. The Q-less QR +// paper states b = 256 for all experiments and its peak-memory claims +// (O(mb + n^2)) assume the tall intermediate is processed in blocks; drivers +// therefore default to this value rather than 0 (unblocked). Callers can still +// set block_size = 0 explicitly to materialize the full m x n intermediate. +inline constexpr int64_t kDefaultGramBlockSize = 256; + +// Env-gated (read once): the sCholQR3 first-pass shift defaults to the paper's +// prescription s = 11*eps*n*trace(G) (FukayaEtAl2020, c = 11). +// RANDLAPACK_SCHOLQR3_SHIFT=eps selects the smaller legacy shift s = eps*trace(G) +// instead (an empirical variant kept for A/B campaigns); the historical value +// "theory" is accepted and means the default. Shared here so the linop and dense +// families read the same knob the same way. +// The static cache means the value is fixed for the process's whole lifetime +// after the first read, so an in-process gtest that flips the env var mid-run +// cannot observe the change; validate this knob via benchmarks, not gtests. +inline bool scholqr3_eps_shift() { + static const bool v = []() { + const char* s = std::getenv("RANDLAPACK_SCHOLQR3_SHIFT"); + return s != nullptr && std::string(s) == "eps"; + }(); + return v; +} + +// Env-gated (read once): the pre-factorization Gram symmetrization +// G <- (G + G^T)/2 defaults ON (the paper's implementation section prescribes +// it, and its roundoff assumption presumes it). RANDLAPACK_CHOL_SYMMETRIZE=0 +// disables it, so potrf factorizes the upper triangle as computed — the +// pre-audit (pre-B5) behavior, kept for A/B campaigns isolating the +// symmetrization's ULP-level effect on borderline pivots. Any other value, +// or unset, means the default. Same static-cache caveat as above: validate +// via benchmarks, not gtests. +inline bool chol_symmetrize() { + static const bool v = []() { + const char* s = std::getenv("RANDLAPACK_CHOL_SYMMETRIZE"); + return s == nullptr || std::string(s) != "0"; + }(); + return v; +} + + +// ============================================================================ +// blocked_preconditioned_gram +// ============================================================================ +// +// Forms the (optionally preconditioned) Gram matrix that cholqr_primitive then +// Cholesky-factorizes. The operator A is matrix-free (only A * B and A^T * B are +// available), so the Gram is built one column block at a time: +// +// R_pre != nullptr : G = R_pre^T (A^T A) R_pre (preconditioned Gram) +// R_pre == nullptr : G = A^T A (plain Gram) +// +// This is the only place A is touched, so it dominates the cost (2 linop applies +// per block); keeping it blocked bounds the scratch to O(n^2 + (m+n) b_eff) +// instead of materializing A (m*n). +// +// A_temp m x b_eff : holds A * B_in for the current block. +// Z_buf n x b_eff : holds A^T * A_temp (used only on the R_pre != nullptr, +// non-skip path; otherwise A^T * A_temp goes straight to G). +// When R_pre == nullptr, an n x n identity is allocated internally so each block +// B_in is a column block of I (i.e. we apply A to the identity to read its columns). +// +// skip_left_factor (R_pre != nullptr only): write A^T A R_pre into G and let the +// caller apply the left R_pre^T factor afterwards (a single TRSM with P) instead +// of a per-block GEMM. Timing accumulators are outputs; ignored when timing==false. +template +void blocked_preconditioned_gram( + GLO& A, + const T* R_pre, + T* G, + int64_t m, int64_t n, int64_t b_eff, + T* A_temp, + T* Z_buf, + bool skip_left_factor, + long& fwd_us, long& adj_us, long& gemm_us, + bool timing) +{ + using std::chrono::steady_clock; + using std::chrono::duration_cast; + using std::chrono::microseconds; + steady_clock::time_point t0, t1; + long fwd_accum = 0, adj_accum = 0, gemm_accum = 0; + + // Unpreconditioned path reads A's columns by applying A to column blocks of + // the identity. We keep a single n x b_eff scratch block (rather than a full + // n x n identity via util::eye, which would cost n^2 memory at scale) and set + // its b_j shifted-diagonal ones each iteration, clearing them again after use. + T* I_block = nullptr; + if (R_pre == nullptr) I_block = new T[n * b_eff](); + + for (int64_t j = 0; j < n; j += b_eff) { + int64_t b_j = std::min(b_eff, n - j); + + // B_in is the j-th column block of R_pre (preconditioned) or of I_n. + const T* B_in; + if (R_pre != nullptr) { + B_in = R_pre + j * n; + } else { + lapack::laset(MatrixType::General, b_j, b_j, (T)0, (T)1, I_block + j, n); + B_in = I_block; + } + + // A_temp = A * B_in (m x b_j) + if (timing) t0 = steady_clock::now(); + A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, m, b_j, n, (T)1.0, B_in, n, (T)0.0, A_temp, m); + if (timing) { t1 = steady_clock::now(); fwd_accum += duration_cast(t1 - t0).count(); } + + if (R_pre && !skip_left_factor) { + // Z_buf = A^T * A_temp ; then G[:, blk] = R_pre^T * Z_buf. + if (timing) t0 = steady_clock::now(); + A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, n, b_j, m, (T)1.0, A_temp, m, (T)0.0, Z_buf, n); + if (timing) { t1 = steady_clock::now(); adj_accum += duration_cast(t1 - t0).count(); } + + if (timing) t0 = steady_clock::now(); + blas::gemm(Layout::ColMajor, Op::Trans, Op::NoTrans, n, b_j, n, (T)1.0, R_pre, n, Z_buf, n, (T)0.0, G + j * n, n); + if (timing) { t1 = steady_clock::now(); gemm_accum += duration_cast(t1 - t0).count(); } + } else { + // skip_left_factor (preconditioned) and the unpreconditioned path both + // write A^T * A_temp straight into G[:, blk]; any left factor is applied + // once after the loop by the caller. + if (timing) t0 = steady_clock::now(); + A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, n, b_j, m, (T)1.0, A_temp, m, (T)0.0, G + j * n, n); + if (timing) { t1 = steady_clock::now(); adj_accum += duration_cast(t1 - t0).count(); } + + // Clear this block's identity ones before the next iteration. + if (R_pre == nullptr) + lapack::laset(MatrixType::General, b_j, b_j, (T)0, (T)0, I_block + j, n); + } + } + + if (I_block) delete[] I_block; + + if (timing) { + fwd_us = fwd_accum; + adj_us = adj_accum; + gemm_us = gemm_accum; + } +} + + +// Materialize Q = A * R^{-1} (m x n) block-by-block via the linop, for the +// test/verify paths of the CholQR-family drivers. R is n x n upper-triangular +// (ld = ldr); Q_out (m x n, leading dimension ldq) is caller-allocated. Forms +// R^{-1} once (n x n trsm), then applies A to its column blocks. Not on any +// timed/algorithmic path. +template +void materialize_Q_from_R(GLO& A, const T* R, int64_t ldr, + int64_t m, int64_t n, int64_t b_eff, T* Q_out, int64_t ldq) { + T* R_inv = new T[n * n]; + RandLAPACK::util::eye(n, n, R_inv); + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, Diag::NonUnit, n, n, T(1), R, ldr, R_inv, n); + for (int64_t j = 0; j < n; j += b_eff) { + int64_t b_j = std::min(b_eff, n - j); + A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, m, b_j, n, T(1), R_inv + j * n, n, T(0), Q_out + j * ldq, ldq); + } + delete[] R_inv; +} + + +// ============================================================================ +// cholqr_primitive: Q-less (preconditioned) Cholesky QR +// ============================================================================ +// +// Computes upper-triangular R such that A * R^{-1} has orthonormal columns. A +// single primitive covers both the plain and preconditioned variants via the +// optional preconditioner P: +// +// P == nullptr (unpreconditioned CholQR): +// G = A^T A; G = R^T R (Cholesky); output R. +// P != nullptr (preconditioned, P upper-triangular): +// R_pre = invert(P) via 'method'; G = R_pre^T A^T A R_pre; +// G = (R^chol)^T R^chol (Cholesky); output R = R^chol * P. +// +// CholQR2 / sCholQR3 chain this: pass the previous iterate's R as P so the +// returned R is the accumulated factor R_k ... R_1. +// +// Adaptive shift: before Cholesky a diagonal shift s = shift_factor * trace(G) is +// added (s = 0 when shift_factor = 0). On a non-PD pivot the shift is grown by +// shift_growth and potrf retried, up to max_retries times; max_retries < 0 means +// unbounded (retry until PD; geometric growth guarantees termination once the +// shift reaches trace(G), where the Gram is diagonally dominant). The Gram is +// computed once; a failed attempt is undone in O(n^2) from the Gram's own strict +// lower triangle plus an O(n) diagonal snapshot, so retries never cost another +// O(m n^2) Gram build and need no n x n backup. +// +// CONTRACT CAVEAT: when a nonzero shift ends up applied (shift_factor > 0, or the +// retry fired), the returned R is the Cholesky factor of G + s I, NOT of G, so +// A R^{-1} is only near-orthonormal: directions of A with singular value below +// sqrt(s) are damped rather than normalized, and R acts as a shifted-CholeskyQR +// preconditioner (Fukaya et al.) rather than a QR factor. Callers that need a true +// factor must run a corrective unshifted pass (CholQR2 / sCholQR3 do). n_retries, +// applied_shift, and gram_trace exist so the caller can detect and report this. +// +// Caller-owned scratch: R_pre (n x n, preconditioned only), G (n x n), A_temp +// (m x b_eff), Z_buf (n x b_eff, preconditioned non-skip only). state is the RNG +// state for the BQRRP method (nullptr otherwise). Timing args are outputs. +// +// Returns 0 on success; 1 on a diag-zero/singular preconditioner; potrf's info on +// a Cholesky breakdown that survived all retries; -1 when a shift or trace is +// non-finite; -2 on invalid shift input (negative shift_factor or shift_growth +// <= 1). m >= n, ldr >= n, and R != nullptr are enforced by randlapack_require +// (throw) rather than a sentinel return, since these are caller bugs, not +// runtime conditions. +template +int cholqr_primitive( + GLO& A, + const T* P, + T* R, int64_t ldr, + PCholQRPrecondMethod method, + int64_t block_size, + T bqrrp_block_ratio, + T* R_pre, + T* G, + T* A_temp, + T* Z_buf, + RandBLAS::RNGState* state, + long& precond_inv_us, + long& fwd_us, long& adj_us, long& gemm_us, long& chol_us, long& update_us, + bool timing, + T shift_factor = T(0), + int max_retries = 0, + T shift_growth = T(10), + int* n_retries = nullptr, // out: number of shift retries used (0 = clean first attempt) + T* applied_shift = nullptr, // out: absolute diagonal shift s in the last attempt (0 = unshifted) + T* gram_trace = nullptr) // out: trace(G), the scale the shift multiplies +{ + using std::chrono::steady_clock; + using std::chrono::duration_cast; + using std::chrono::microseconds; + int64_t m = A.n_rows; + int64_t n = A.n_cols; + int64_t b_eff = (block_size > 0 && block_size < n) ? block_size : n; + const bool preconditioned = (P != nullptr); + + // Basic input validation, shared by every driver that reaches this + // primitive: cholqr_iterate (CholQR/CholQR2/sCholQR3, via either overload + // below) and CQRRT_linops (which calls this overload directly, bypassing + // cholqr_iterate). Caller bugs, so throw rather than return a sentinel. + randlapack_require(m >= n) << "cholqr_primitive: operator must be tall (m=" << m << " < n=" << n << ")"; + randlapack_require(ldr >= n) << "cholqr_primitive: ldr=" << ldr << " < n=" << n; + randlapack_require(R != nullptr) << "cholqr_primitive: R buffer is null"; + + // Reset every out-param up front: the early-return failure paths below must + // not leave a previous call's values behind (a stale n_retries was being + // re-summed by cholqr_iterate's per-pass accumulation). + if (n_retries) *n_retries = 0; + if (applied_shift) *applied_shift = T(0); + if (gram_trace) *gram_trace = T(0); + if (shift_factor < T(0)) { + std::fprintf(stderr, + "[cholqr_primitive] FAIL: negative shift_factor (%g) is invalid\n", + (double)shift_factor); + return -2; + } + if (shift_growth <= T(1)) { + std::fprintf(stderr, + "[cholqr_primitive] FAIL: shift_growth (%g) must be > 1 (<= 1 defeats " + "the geometric-growth retry termination argument)\n", + (double)shift_growth); + return -2; + } + + steady_clock::time_point t0, t1; + if (timing) t0 = steady_clock::now(); + + // ---- Step 1: R_pre = invert(P) (preconditioned only) ---- + if (preconditioned) { + switch (method) { + case PCholQRPrecondMethod::TRSM_IDENTITY: { + if (!RandLAPACK::util::diag_is_nonzero(n, P, n)) { + std::fprintf(stderr, "[cholqr_primitive] FAIL: TRSM_IDENTITY diag_is_nonzero(P) failed (P has ~0 diagonal entry)\n"); + return 1; + } + RandLAPACK::util::eye(n, n, R_pre); + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, Diag::NonUnit, n, n, T(1), P, n, R_pre, n); + if (n > 1) + lapack::laset(MatrixType::Lower, n - 1, n - 1, T(0), T(0), R_pre + 1, n); + break; + } + case PCholQRPrecondMethod::TRTRI: { + if (!RandLAPACK::util::diag_is_nonzero(n, P, n)) { + std::fprintf(stderr, "[cholqr_primitive] FAIL: TRTRI diag_is_nonzero(P) failed (P has ~0 diagonal entry)\n"); + return 1; + } + lapack::lacpy(MatrixType::Upper, n, n, P, n, R_pre, n); + if (n > 1) + lapack::laset(MatrixType::Lower, n - 1, n - 1, T(0), T(0), R_pre + 1, n); + int trtri_info = lapack::trtri(Uplo::Upper, Diag::NonUnit, n, R_pre, n); + if (trtri_info) { + std::fprintf(stderr, "[cholqr_primitive] FAIL: lapack::trtri returned info=%d\n", trtri_info); + return 1; + } + break; + } + case PCholQRPrecondMethod::GEQP3: + case PCholQRPrecondMethod::BQRRP: { + #if defined(__APPLE__) + // The whole QRCP-based preconditioner path (GEQP3/BQRRP, then ungqr + + // lapmr to form Pi R_tri^{-1} Q^T) pulls in LAPACK / BQRRP routines that + // are unsupported under Apple Accelerate; the sibling rl_cqrrpt.hh / + // rl_hqrrp.hh guard the whole QRCP path the same way. The standard + // CholQR / CholQR2 / sCholQR3 methods use TRSM_IDENTITY and never reach + // here, so only the stabilized QRCP preconditioner is disabled on macOS. + (void)bqrrp_block_ratio; (void)state; + std::fprintf(stderr, "[cholqr_primitive] FAIL: GEQP3/BQRRP preconditioning is unsupported on Apple Accelerate.\n"); + return 1; + #else + // Invert P stably via column-pivoted QR. Column pivoting gives + // P Pi = Q R_tri (Pi the pivot permutation), + // so P^{-1} = Pi R_tri^{-1} Q^T. We build that below from the QRCP + // outputs; this avoids the kappa(P) error amplification a direct + // triangular solve against an ill-conditioned P would incur. + T* P_copy = new T[n * n](); + lapack::lacpy(MatrixType::Upper, n, n, P, n, P_copy, n); // P_copy = P (upper); lower stays 0 + if (!RandLAPACK::util::diag_is_nonzero(n, P_copy, n)) { + std::fprintf(stderr, "[cholqr_primitive] FAIL: GEQP3/BQRRP diag_is_nonzero(P) failed\n"); + delete[] P_copy; return 1; + } + + int64_t* jpiv = new int64_t[n](); + T* tau_qr = new T[n]; + + if (method == PCholQRPrecondMethod::GEQP3) { + lapack::geqp3(n, n, P_copy, n, jpiv, tau_qr); + } else { + if (state == nullptr) { + std::fprintf(stderr, "[cholqr_primitive] FAIL: BQRRP called with state=nullptr\n"); + delete[] P_copy; delete[] jpiv; delete[] tau_qr; + return 1; + } + T ratio = bqrrp_block_ratio; + if (ratio == T(1.0)) { + if (n <= 2000) ratio = T(1.0); + else if (n <= 8000) ratio = T(0.5); + else ratio = T(1.0) / T(32); + } + int64_t bqrrp_b = std::max(int64_t(1), (int64_t)(n * ratio)); + RandLAPACK::BQRRP bqrrp(false, bqrrp_b); + bqrrp.call(n, n, P_copy, n, T(1), tau_qr, jpiv, *state); + } + + // After QRCP, P_copy holds R_tri (upper) + Householder vectors (lower). + // Pull out R_tri, rebuild Q in place, then form R_tri^{-1} Q^T. + T* R_buf = new T[n * n](); + lapack::lacpy(MatrixType::Upper, n, n, P_copy, n, R_buf, n); // R_buf = R_tri + lapack::ungqr(n, n, n, P_copy, n, tau_qr); // P_copy = Q + RandLAPACK::util::transpose_square(P_copy, n); // P_copy = Q^T + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, Diag::NonUnit, n, n, T(1), R_buf, n, P_copy, n); // P_copy = R_tri^{-1} Q^T + + // R_pre = Pi (R_tri^{-1} Q^T): jpiv is the column-pivot permutation, applied to + // the ROWS here, so copy R^{-1} Q^T over and apply it with lapmr (row permute). + lapack::lacpy(MatrixType::General, n, n, P_copy, n, R_pre, n); + lapack::lapmr(false, n, n, R_pre, n, jpiv); + + delete[] P_copy; + delete[] R_buf; + delete[] jpiv; + delete[] tau_qr; + break; + #endif + } + } + } + + if (timing) { + t1 = steady_clock::now(); + precond_inv_us = preconditioned ? duration_cast(t1 - t0).count() : 0; + } + + // ---- Step 2: form the Gram G ---- + // Preconditioned TRSM_IDENTITY/TRTRI defer the left R_pre^T factor to a single + // TRSM (cheaper, O(n^3/2), and equally stable since P is preserved); GEQP3/BQRRP + // keep the explicit per-block GEMM with R_pre^T to preserve the QRCP-accurate + // R_pre (a TRSM with the ill-conditioned P would re-amplify the error). The + // unpreconditioned path forms plain A^T A. + // + // RANDLAPACK_GRAM_LEFT=gemm forces the per-block GEMM with R_pre^T even for + // TRSM_IDENTITY/TRTRI: the paper's stability analysis models the + // explicit-multiply path, the default runs the TRSM path; this knob + // lets a campaign measure both. Only the gemm direction can be forced: making + // GEQP3/BQRRP use the TRSM would reintroduce the error their construction avoids. + // Static-cached: read once per process, so an in-process gtest toggling the + // env var mid-run cannot observe the change; validate via benchmarks, not gtests. + static const bool force_gemm_left = []() { + const char* s = std::getenv("RANDLAPACK_GRAM_LEFT"); + return s != nullptr && std::string(s) == "gemm"; + }(); + const bool use_trsm_at_end = preconditioned + && !force_gemm_left + && (method == PCholQRPrecondMethod::TRSM_IDENTITY + || method == PCholQRPrecondMethod::TRTRI); + + blocked_preconditioned_gram(A, preconditioned ? R_pre : (const T*)nullptr, G, m, n, b_eff, A_temp, Z_buf, /*skip_left_factor=*/use_trsm_at_end, fwd_us, adj_us, gemm_us, timing); + + if (use_trsm_at_end) { + // G := P^{-T} G (= R_pre^T A^T A R_pre, since R_pre = P^{-1}). + if (timing) t0 = steady_clock::now(); + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::Trans, Diag::NonUnit, n, n, T(1), P, n, G, n); + if (timing) { t1 = steady_clock::now(); gemm_us += duration_cast(t1 - t0).count(); } + } + + // Symmetrize G <- (G + G^T)/2 before factorizing, as the paper's + // implementation section states and its roundoff assumption presumes. The + // two triangles come from independent floating-point summations, so they + // agree only up to formation rounding; the average is written into BOTH + // triangles so the retry-restore below reproduces the symmetrized matrix + // exactly. O(n^2), negligible next to potrf's n^3/3. + // RANDLAPACK_CHOL_SYMMETRIZE=0 (chol_symmetrize() above) skips this, so + // potrf sees the upper triangle as computed. In that mode a retry restores + // the upper from the as-computed strict LOWER, i.e. the transpose of what + // attempt 0 factorized — the two differ by formation rounding only, which + // is immaterial under a shift >= eps*trace (same argument as the restore + // note below). + if (chol_symmetrize()) { + if (timing) t0 = steady_clock::now(); + for (int64_t j = 0; j < n; ++j) { + for (int64_t i = 0; i < j; ++i) { + T avg = (G[i + j * n] + G[j + i * n]) / T(2); + G[i + j * n] = avg; + G[j + i * n] = avg; + } + } + if (timing) { t1 = steady_clock::now(); gemm_us += duration_cast(t1 - t0).count(); } + } + + // ---- Step 3: Cholesky G = (R^chol)^T R^chol, with adaptive-shift retry ---- + // + // The retry block exists because the (preconditioned) Gram can pick up a + // non-PD pivot from rounding, amplified by kappa(R_pre)^2 in the iter-2/3 + // Gram of CholQR2/sCholQR3. On each potrf failure the Gram is restored, the + // diagonal shift grown, and potrf retried, so retries stay O(n^2). Seeding + // from eps on the first bump lets an unshifted (shift_factor==0) caller keep a + // clean first attempt while still being rescued if the Gram is non-PD. + // + // Restore mechanics: potrf(Upper) reads and overwrites only the upper triangle + // plus diagonal, and after the symmetrization above the strict lower triangle + // holds exactly the symmetrized values, survives a failed attempt untouched, + // and serves as the restore source for the upper. Only the diagonal + // needs an O(n) snapshot. Nothing downstream reads G's strict lower (potrf, + // the output lacpy, and the trmm are all Upper), so it stays unzeroed. + T* diag_backup = (max_retries != 0) ? new T[n] : nullptr; + T trace_G = 0; + for (int64_t i = 0; i < n; ++i) { + T d = G[i * (n + 1)]; + trace_G += d; + if (diag_backup) diag_backup[i] = d; + } + if (gram_trace) *gram_trace = trace_G; + + if (timing) t0 = steady_clock::now(); + + int info = 0; + int attempt = 0; + T current_shift_factor = shift_factor; + T last_shift = T(0); // absolute shift present in G on the most recent attempt + // max_retries < 0 means "unbounded", which with a geometrically growing shift is + // *usually* fine: the shift eventually makes G diagonally dominant and potrf + // succeeds. But the argument fails on non-finite data (an inf/NaN in G, or a shift + // that overflows to inf, gives a Gram potrf can never factor, and the loop then never + // terminates. kUnboundedRetryCeiling is a backstop for exactly that case: it is far + // above any legitimate retry count (each attempt multiplies the shift by + // shift_growth, so tens of attempts already span the entire exponent range), so it + // cannot truncate a run that would otherwise have succeeded. + constexpr int kUnboundedRetryCeiling = 128; + for (; (max_retries < 0) ? (attempt < kUnboundedRetryCeiling) : (attempt <= max_retries); ++attempt) { + if (attempt > 0) { + // Restore the upper triangle from the untouched strict lower plus the + // saved diagonal, and grow the shift (seed at eps if we started at 0). + for (int64_t j = 0; j < n; ++j) { + for (int64_t i = 0; i < j; ++i) G[i + j * n] = G[j + i * n]; + G[j * (n + 1)] = diag_backup[j]; + } + current_shift_factor = (current_shift_factor > T(0)) + ? current_shift_factor * shift_growth + : std::numeric_limits::epsilon(); + // A non-finite shift can never rescue the factorization; bail out rather + // than spin. Same for a non-finite trace, which poisons every shift below. + if (!std::isfinite(current_shift_factor) || !std::isfinite(trace_G)) { + info = -1; + break; + } + } + if (current_shift_factor > T(0)) { + T shift = current_shift_factor * trace_G; + // The factor and the trace can both be finite while their product + // overflows; an infinite shift can never rescue potrf, so bail. + if (!std::isfinite(shift)) { + info = -1; + break; + } + for (int64_t i = 0; i < n; ++i) G[i * (n + 1)] += shift; + last_shift = shift; + } + info = lapack::potrf(Uplo::Upper, n, G, n); + if (info == 0) break; + } + // Retries actually performed: `attempt` on success (attempt 0 = clean first + // try). When the loop ends without success (exhaustion or a non-finite-shift + // bail), the attempt indexed by `attempt` never ran, so the count is one less. + if (n_retries) *n_retries = (info == 0) ? attempt : (attempt > 0 ? attempt - 1 : 0); + if (applied_shift) *applied_shift = last_shift; + + // A nominally unshifted call (shift_factor == 0) that only succeeded via the + // adaptive retry is no longer the fixed unshifted algorithm the caller named: + // R factors G + sI, which acts as an extra regularizing preconditioner. Say + // so loudly, since result rows keep the plain algorithm name and only the + // chol_retries / chol_shift columns reveal the rescue. + if (info == 0 && attempt > 0 && shift_factor == T(0)) { + std::fprintf(stderr, + "[cholqr_primitive] NOTE: nominally unshifted Cholesky rescued by the " + "adaptive shift after %d retry(ies) (shift=%.3e, shift/trace(G)=%.3e); " + "this run measures the adaptive-shift variant of the calling algorithm\n", + attempt, (double)last_shift, + (double)(trace_G > T(0) ? last_shift / trace_G : T(0))); + } + + if (info) { + // Report the retries actually made, not the configured limit: + // with max_retries = -1 the old message printed "-1 retries". + // info == -1 is our own non-finite-shift/trace bail (Step 3 above); + // potrf was never called on that attempt, so it did not return -1. + // Any other nonzero info is potrf's own Cholesky-breakdown code. + if (info == -1) { + std::fprintf(stderr, + "[cholqr_primitive] FAIL: non-finite shift or Gram trace after %d " + "attempt(s) (final shift_factor=%g); potrf was not called\n", + attempt, (double)current_shift_factor); + } else { + std::fprintf(stderr, + "[cholqr_primitive] FAIL: lapack::potrf returned info=%d after %d " + "attempt(s) (final shift_factor=%g)\n", + info, attempt, (double)current_shift_factor); + } + delete[] diag_backup; + return info; + } + delete[] diag_backup; + + if (timing) { t1 = steady_clock::now(); chol_us = duration_cast(t1 - t0).count(); } + + // ---- Step 4: output R ---- + // Unpreconditioned: R = R^chol. Preconditioned: R = R^chol * P (accumulates + // the running factor), computed in place by seeding R with P then a TRMM. + if (timing) t0 = steady_clock::now(); + if (preconditioned) { + lapack::lacpy(MatrixType::Upper, n, n, P, n, R, ldr); + if (n > 1) + lapack::laset(MatrixType::Lower, n - 1, n - 1, T(0), T(0), R + 1, ldr); + blas::trmm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, Diag::NonUnit, n, n, T(1), G, n, R, ldr); + } else { + lapack::lacpy(MatrixType::Upper, n, n, G, n, R, ldr); + if (n > 1) + lapack::laset(MatrixType::Lower, n - 1, n - 1, T(0), T(0), R + 1, ldr); + } + if (timing) { t1 = steady_clock::now(); update_us = preconditioned ? duration_cast(t1 - t0).count() : 0; } + + return 0; +} + + +// Unpreconditioned convenience overload: plain CholQR (P = nullptr). Owns the G / +// A_temp scratch the general primitive needs and forwards. This is the entry point +// for CholQR and for iteration 1 of CholQR2 / sCholQR3. +template +int cholqr_primitive( + GLO& A, + T* R, int64_t ldr, + T shift_factor, + int64_t block_size, + long& fwd_us, long& adj_us, long& chol_us, + bool timing, + int max_retries = 0, + T shift_growth = T(10), + int* n_retries = nullptr, + T* applied_shift = nullptr, + T* gram_trace = nullptr) +{ + int64_t m = A.n_rows; + int64_t n = A.n_cols; + int64_t b_eff = (block_size > 0 && block_size < n) ? block_size : n; + + // Validated again here (duplicating the checks in the general overload + // below) so a bad m/ldr/R throws before G/A_temp are allocated, rather + // than leaking them when the general overload's own check throws. + randlapack_require(m >= n) << "cholqr_primitive: operator must be tall (m=" << m << " < n=" << n << ")"; + randlapack_require(ldr >= n) << "cholqr_primitive: ldr=" << ldr << " < n=" << n; + randlapack_require(R != nullptr) << "cholqr_primitive: R buffer is null"; + + T* G = new T[n * n](); + T* A_temp = new T[m * b_eff]; + + long precond_inv_us = 0, gemm_us = 0, update_us = 0; + int info = cholqr_primitive( + A, /*P=*/(const T*)nullptr, R, ldr, + PCholQRPrecondMethod::TRSM_IDENTITY, // ignored when P == nullptr + block_size, /*bqrrp_block_ratio=*/T(1), /*R_pre=*/(T*)nullptr, + G, A_temp, /*Z_buf=*/(T*)nullptr, + /*state=*/(RandBLAS::RNGState*)nullptr, + precond_inv_us, fwd_us, adj_us, gemm_us, chol_us, update_us, + timing, shift_factor, max_retries, shift_growth, n_retries, + applied_shift, gram_trace); + + delete[] G; + delete[] A_temp; + return info; +} + + +// ============================================================================ +// cholqr_iterate: the shared CholQR-family engine +// ============================================================================ +// +// Runs `num_iters` CholQR passes and returns the accumulated R (= R_k ... R_1): +// iter 1 : unpreconditioned CholQR with shift_iter1. +// iter 2..num_iters : preconditioned CholQR with the previous iterate as P +// (TRSM_IDENTITY) and shift_iter_rest. +// This is the single engine behind CholQR (num_iters = 1), CholQR2 (2), and +// sCholQR3 (3); the only differences are the pass count and the per-pass shift +// (CholQR/CholQR2 start unshifted, shift_iter1 = shift_iter_rest = 0; sCholQR3 +// uses eps). The adaptive-shift retry inside cholqr_primitive handles potrf +// breakdown; max_retries < 0 means unbounded. +// +// Scratch for the preconditioned passes (G, R_pre, P_prev, A_temp, Z_buf) is owned +// internally and only allocated when num_iters > 1, so the num_iters = 1 (plain +// CholQR) path keeps its lean footprint. +// +// iter_times (optional, length 5*num_iters): per pass [fwd, adj, gemm, chol, upd]; +// gemm and upd are 0 for the unpreconditioned first pass. +// +// applied_shifts / gram_traces (optional, length num_iters): per pass, the absolute +// diagonal shift the successful potrf attempt carried (0 = unshifted) and the trace +// of that pass's Gram. A nonzero pass-1 shift means the accumulated R factors +// G + s I rather than G (see the contract caveat on cholqr_primitive); entries for +// passes not reached stay 0. +// +// Returns 0 on success, or the 1-based index of the pass whose factorization +// failed. That failure can have any of cholqr_primitive's causes: retry +// exhaustion, a singular preconditioner, a non-finite shift, or invalid input; +// the specific cause is printed to stderr by cholqr_primitive, not encoded in +// this return value. +template +int cholqr_iterate( + GLO& A, T* R, int64_t ldr, int64_t block_size, + int num_iters, T shift_iter1, T shift_iter_rest, + int max_retries, T shift_growth, bool timing, + long* iter_times = nullptr, + int* n_retries_total = nullptr, // out: total shift retries summed across all passes + T* applied_shifts = nullptr, // out, length num_iters: absolute shift per pass + T* gram_traces = nullptr) // out, length num_iters: trace of each pass's Gram +{ + int64_t m = A.n_rows; + int64_t n = A.n_cols; + int64_t b_eff = (block_size > 0 && block_size < n) ? block_size : n; + int total_retries = 0, pass_retries = 0; + if (applied_shifts) std::fill(applied_shifts, applied_shifts + num_iters, T(0)); + if (gram_traces) std::fill(gram_traces, gram_traces + num_iters, T(0)); + auto rec = [&](int it, long fwd, long adj, long gemm, long chol, long upd) { + if (iter_times) { + long* t = iter_times + 5 * (it - 1); + t[0] = fwd; t[1] = adj; t[2] = gemm; t[3] = chol; t[4] = upd; + } + }; + + // ---- Iter 1: unpreconditioned CholQR ---- + long fwd1 = 0, adj1 = 0, chol1 = 0; + int info = cholqr_primitive(A, R, ldr, shift_iter1, block_size, fwd1, adj1, chol1, timing, max_retries, shift_growth, &pass_retries, + applied_shifts ? &applied_shifts[0] : nullptr, + gram_traces ? &gram_traces[0] : nullptr); + total_retries += pass_retries; + if (info != 0) { if (n_retries_total) *n_retries_total = total_retries; return 1; } + rec(1, fwd1, adj1, 0, chol1, 0); + if (num_iters <= 1) { if (n_retries_total) *n_retries_total = total_retries; return 0; } + + // ---- Iters 2..num_iters: preconditioned CholQR with P = previous R ---- + T* G = new T[n * n](); + T* R_pre = new T[n * n](); + T* P_prev = new T[n * n](); + T* A_temp = new T[m * b_eff]; + T* Z_buf = new T[n * b_eff]; + for (int it = 2; it <= num_iters; ++it) { + lapack::lacpy(MatrixType::Upper, n, n, R, ldr, P_prev, n); + if (n > 1) lapack::laset(MatrixType::Lower, n - 1, n - 1, T(0), T(0), P_prev + 1, n); + long precond_inv = 0, fwd = 0, adj = 0, gemm = 0, chol = 0, upd = 0; + info = cholqr_primitive( + A, P_prev, R, ldr, PCholQRPrecondMethod::TRSM_IDENTITY, + block_size, /*bqrrp_block_ratio=*/T(1), R_pre, G, A_temp, Z_buf, + /*state=*/(RandBLAS::RNGState*)nullptr, + precond_inv, fwd, adj, gemm, chol, upd, timing, + shift_iter_rest, max_retries, shift_growth, &pass_retries, + applied_shifts ? &applied_shifts[it - 1] : nullptr, + gram_traces ? &gram_traces[it - 1] : nullptr); + total_retries += pass_retries; + if (info != 0) { + delete[] G; delete[] R_pre; delete[] P_prev; delete[] A_temp; delete[] Z_buf; + if (n_retries_total) *n_retries_total = total_retries; + return it; + } + rec(it, fwd, adj, gemm, chol, precond_inv + upd); + } + delete[] G; delete[] R_pre; delete[] P_prev; delete[] A_temp; delete[] Z_buf; + if (n_retries_total) *n_retries_total = total_retries; + return 0; +} + +} // namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_blendenpik.hh b/RandLAPACK/drivers/rl_blendenpik.hh new file mode 100644 index 000000000..7be572fe5 --- /dev/null +++ b/RandLAPACK/drivers/rl_blendenpik.hh @@ -0,0 +1,302 @@ +#pragma once + +// Public API: Blendenpik_linops: sketch-and-precondition least-squares solver. +// +// Classical Blendenpik (Avron, Maymounkov, Toledo 2010) for min ||b - A x||_2 on +// a tall LinearOperator A (m x n, m >= n): +// 1. sketch Ask = S A (S a d x m sparse SASO map, d = d_factor * n) +// 2. unpivoted Householder QR of the sketch: [~, R] = qr(Ask) +// 3. R is a right preconditioner: A R^{-1} is nearly orthonormal (kappa ~ 1) +// 4. solve min ||b - (A R^{-1}) y|| by matrix-free LSQR; return x = R^{-1} y. +// +// This is the sparse-projection variant (SASO instead of Blendenpik's SRFT). It is +// an INDEPENDENT solver: no mu-regularization, no iterative refinement. +// Reference: Algorithm SPO1 in Avron, Maymounkov, Toledo (2010). +// +// Sketch-and-solve initialization (`warm_start`, default ON): +// LSQR is started from x0 = R^{-1}(Q^T(S b)), the solution of the sketched problem, +// instead of from zero. Epperly, Meier and Nakatsukasa (arXiv:2406.03468v3, sec. 3.1) +// note this initialization is "necessary for the method to be forward stable" and that +// it is an optional setting in the original Blendenpik code; starting from zero is a +// documented cause of stagnating short of the attainable accuracy. Because rl_lsqr +// always starts from zero internally, the warm start is applied here as an equivalent +// shift: solve for the correction dx against the residual r0 = b - A x0, then return +// x = x0 + dx. That leaves rl_lsqr untouched, so the five Q-less QR methods that share +// it are provably unaffected. +// +// NOTE this remains one-shot sketch-and-precondition (no iterative refinement), which +// the same reference proves is not backward stable. Residual stagnation ABOVE the +// backward-stable level is therefore expected behaviour, not a defect. + +#include "rl_util.hh" +#include "rl_blaspp.hh" +#include "rl_blas2_threads.hh" +#include "rl_exceptions.hh" +#include "rl_lapackpp.hh" +#include "rl_lsqr.hh" +#include "../linops/rl_concepts.hh" + +#include +#include +#include +#include +#include +#include + +namespace RandLAPACK { + + +/// @brief Blendenpik (sparse-sketch + QR preconditioner + LSQR) for tall LS. +template +class Blendenpik_linops { + public: + bool timing; + T tol; ///< LSQR stopping tolerance (atol = btol = tol). + int max_iters; ///< LSQR iteration cap. + int64_t nnz; ///< SASO nonzeros per column (sparse projection). + int lsqr_iters; ///< LSQR iterations used on the last call (output). + /// Start LSQR from the sketch-and-solve solution rather than from zero. + /// Required for forward stability (see the file header); on by default. + bool warm_start; + /// Stop after the sketch-and-solve initial guess and return it as x, skipping + /// LSQR entirely (implies warm_start). Lets a caller reuse this class as a + /// standalone sketch-and-solve solver, e.g. to warm-start IterRefineLSQ with + /// the exact same x0 Blendenpik uses, isolating the initialization effect. + bool init_only; + + // [0]=sketch, [1]=qr, [2]=lsqr, [3]=total, [4]=x0 setup (microseconds). + // Slot 4 is the sketch-and-solve x0 build (ormqr + trsv + one operator + // apply for r0); it is not included in slot [3]. + std::vector times; + + Blendenpik_linops(bool time_subroutines, T ep) { + timing = time_subroutines; + tol = (ep > (T)0) ? ep : std::numeric_limits::epsilon(); + max_iters = 0; // callers MUST set this before call(); see the require below + nnz = 4; // sparse projection, 4 nnz/col (benchmark CLI overrides; + // NOT the CQRRT default, which is 2) + lsqr_iters = 0; + warm_start = true; + init_only = false; + } + + ~Blendenpik_linops() { + delete[] R_out; + } + + // R_out is an owned raw pointer with no callers copying this object + // (only constructed, called, and read); disable copy so a copy never + // silently double-frees or aliases the buffer. + Blendenpik_linops(const Blendenpik_linops&) = delete; + Blendenpik_linops& operator=(const Blendenpik_linops&) = delete; + + /// Solve min ||b - A x||_2. A is m x n; b length m; x length n (output). + /// d_factor sets the sketch size d = d_factor * n (>= 1, typ. 4). + template + int call(GLO& A, const T* b, int64_t m, T* x, int64_t n, + T d_factor, RandBLAS::RNGState& state) + { + // Reset every output member up front so an early-return path + // (including the rank-deficient sketch guard below) never leaves + // a previous call's values on a reused object. + converged = false; + final_relres = (T)-1; + lsqr_iters = 0; + lsqr_stop_test = 0; + lsqr_op_times.clear(); + times.clear(); + delete[] R_out; + R_out = nullptr; + R_out_sz = 0; + + // Hoisted ahead of every allocation below (was checked only at + // Step 4, after six new[] and the sketch+QR had already run, + // so the throw leaked all of it). + randlapack_require(init_only || max_iters > 0) + << "Blendenpik_linops: set max_iters before call() unless init_only is set (no default cap)"; + + using clock = std::chrono::steady_clock; + using std::chrono::duration_cast; using std::chrono::microseconds; + long t_sketch = 0, t_qr = 0, t_lsqr = 0, t_x0 = 0; + auto total_start = clock::now(); + // Local view only: init_only implies a warm start (x0 IS the output), + // but must not mutate the member, or a reused object would be + // silently warm-started on later calls. + const bool ws = warm_start || init_only; + + int64_t d = (int64_t)(d_factor * (T)n); + if (d < n) d = n; + + // Allocation (and its first-touch zeroing) is timed INTO the sketch + // slot: the Q-less drivers all time their allocations into the + // build total, so leaving Blendenpik's outside every slot would + // make its build bar incomparable. + auto t0 = clock::now(); + T* Ask = new T[d * n](); // sketch (d x n, ColMajor) + T* tau = new T[n](); + T* R = new T[n * n](); // preconditioner (upper triangular) + T* Sb = ws ? new T[d]() : nullptr; // sketched RHS, for x0 + T* x0 = ws ? new T[n]() : nullptr; + T* r0 = ws ? new T[m]() : nullptr; + auto cleanup = [&]() { + delete[] Ask; delete[] tau; delete[] R; + delete[] Sb; delete[] x0; delete[] r0; + }; + + // ---- Step 1: Ask = S A (sparse SASO, applied from the right by the operator) ---- + RandBLAS::SparseDist DS(d, m, this->nnz); + RandBLAS::SparseSkOp S(DS, state); + state = S.next_state; + RandBLAS::fill_sparse(S); + A(blas::Side::Right, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + d, n, m, (T)1.0, S, (T)0.0, Ask, d); + // Sketch the RHS with the SAME S, so x0 solves the sketched LS problem + // min ||Ask x - Sb||. Treat b as an m x 1 matrix. + if (ws) { + RandBLAS::sketch_general(blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + d, 1, m, (T)1.0, S, 0, 0, b, m, (T)0.0, Sb, d); + } + if (timing) t_sketch = duration_cast(clock::now() - t0).count(); + + // ---- Step 2: unpivoted Householder QR of the sketch; R = upper(Ask) ---- + t0 = clock::now(); + lapack::geqrf(d, n, Ask, d, tau); + lapack::lacpy(MatrixType::Upper, n, n, Ask, d, R, n); + if (n > 1) lapack::laset(MatrixType::Lower, n - 1, n - 1, (T)0, (T)0, R + 1, n); + if (timing) t_qr = duration_cast(clock::now() - t0).count(); + + if (!RandLAPACK::util::diag_is_nonzero(n, R, n)) { + std::fprintf(stderr, "[Blendenpik] FAIL: sketch R has a ~0 diagonal (rank-deficient sketch)\n"); + cleanup(); + return 1; + } + + // ---- Step 3: sketch-and-solve initial guess x0 = R^{-1} (Q^T (S b)) ---- + // Q is the implicit factor from geqrf(Ask); apply Q^T with ormqr rather than + // forming it, then one triangular solve. The dominant cost is the one full + // operator apply for r0. Timed into its own slot (times[4]) so + // callers reporting sketch+qr and lsqr do not drop it. + if (ws) { + t0 = clock::now(); + lapack::ormqr(blas::Side::Left, blas::Op::Trans, d, 1, n, Ask, d, tau, Sb, d); + std::copy(Sb, Sb + n, x0); + { Blas2ThreadGuard tg(n); // cap threads: see rl_blas2_threads.hh + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, blas::Op::NoTrans, + blas::Diag::NonUnit, n, R, n, x0, 1); + } + // r0 = b - A x0: LSQR then solves for the correction against this residual. + A(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, x0, n, (T)0.0, r0, m); + for (int64_t i = 0; i < m; ++i) r0[i] = b[i] - r0[i]; + if (timing) t_x0 = duration_cast(clock::now() - t0).count(); + } + + // init_only: the sketch-and-solve x0 IS the answer; skip LSQR. Report the + // true relative residual of x0 so callers can log the warm start's quality. + if (init_only) { + std::copy(x0, x0 + n, x); + lsqr_iters = 0; + converged = false; // no tolerance was pursued + T nb = blas::nrm2(m, b, 1); + final_relres = (nb > (T)0) ? blas::nrm2(m, r0, 1) / nb : (T)-1; + if (timing) { + long total = duration_cast(clock::now() - total_start).count(); + this->times = {t_sketch, t_qr, 0, total, t_x0}; + } + // Hand R to the caller instead of copying it: R_out now owns the + // buffer, so cleanup() below must not free it too. + R_out = R; R_out_sz = n * n; R = nullptr; + cleanup(); + return 0; + } + + // ---- Step 4: LSQR on A with right preconditioner R; x = R^{-1} y ---- + long lt[4] = {0, 0, 0, 0}; + t0 = clock::now(); + int st = 0; + if (ws) { + T nb = blas::nrm2(m, b, 1); + T nr0 = blas::nrm2(m, r0, 1); + if (nr0 <= tol * nb) { + // x0 already meets the caller's tolerance on the TRUE + // residual ||b - A x0|| / ||b||. Running LSQR from here + // would otherwise pursue S1 (||dx-residual|| <= btol * + // ||r0||) which is orders stricter than tol whenever + // ||r0|| << ||b|| (an unintentionally harder target + // than every other method's stop criterion pursues). + std::copy(x0, x0 + n, x); + lsqr_iters = 0; + lsqr_stop_test = 1; // S1 (residual) already satisfied + lsqr_op_times.assign(4, 0L); + // -1 sentinel for ||b|| == 0, matching init_only above: + // the ratio is undefined, not zero. + final_relres = (nb > (T)0) ? nr0 / nb : (T)-1; + converged = true; + } else { + // LSQR solves for the correction dx against r0 = b - A x0. + // btol_eff makes its S1 test on ||r0 - à dx|| / ||r0|| + // equivalent to the caller's ||b - A x|| / ||b|| <= tol, + // since b - A(x0 + dx) = r0 - A dx. + T btol_eff = tol * nb / nr0; + st = RandLAPACK::lsqr(A, m, n, R, n, r0, x, + tol, btol_eff, max_iters, lsqr_iters, lt, + &final_relres, &lsqr_stop_test); + lsqr_op_times.assign(lt, lt + 4); + // Undo the shift: LSQR solved for the correction, so add the + // initial guess back. + blas::axpy(n, (T)1.0, x0, 1, x, 1); + // LSQR normalized its residual by ||r0||, not ||b||. Rescale + // so the reported number means the same thing as for every + // other method (||b - A x|| / ||b||); the numerator is + // already the true residual (same identity as above). + if (final_relres >= (T)0 && nb > (T)0) final_relres *= nr0 / nb; + // LSQR always returns a valid iterate x (best-so-far), so + // hitting the cap is NOT a hard failure for the SOLUTION, + // but it does mean the requested tolerance was not met, and + // callers must be able to see that. + converged = (st == 0); + } + } else { + st = RandLAPACK::lsqr(A, m, n, R, n, b, x, + tol, tol, max_iters, lsqr_iters, lt, + &final_relres, &lsqr_stop_test); + lsqr_op_times.assign(lt, lt + 4); + converged = (st == 0); + } + if (timing) t_lsqr = duration_cast(clock::now() - t0).count(); + + if (timing) { + long total = duration_cast(clock::now() - total_start).count(); + this->times = {t_sketch, t_qr, t_lsqr, total, t_x0}; + } + // Expose the sketch R factor so the caller can report Q = A R^{-1} + // orthogonality. Hand off the buffer (cleanup() must not free it). + R_out = R; R_out_sz = n * n; R = nullptr; + + cleanup(); + return 0; // x is a valid iterate (converged or capped); only the rank-deficient + // sketch guard above returns 1 (hard failure, no usable R). + } + + /// Sketch R factor from the last call (n x n, ColMajor upper-triangular), + /// for Q = A R^{-1} orthogonality reporting. Owned by this object (the + /// destructor frees it); nullptr / 0 when no successful call has run yet. + T* R_out = nullptr; + int64_t R_out_sz = 0; + /// Whether LSQR met its tolerance (false = hit the iteration cap). + bool converged = false; + /// LSQR's own ||b - A x|| / ||b|| at termination. Previously Blendenpik never + /// requested this from lsqr, so its CSV column was structurally -1 while every + /// other method carried a real number. + T final_relres = (T)-1; + /// Which LSQR stopping test ended the run: 1 = S1 (residual), 2 = S2 + /// (normal-equation test; fires at the LS floor), 0 = iteration cap. + int lsqr_stop_test = 0; + /// LSQR's internal operator split [fwd_us, adj_us, trsm_us, total_us] from + /// the last call. If discarded instead, the benchmarks' op-split + /// columns would be forced to -1 sentinels. + std::vector lsqr_op_times; +}; + + +} // namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_cholqr_dense.hh b/RandLAPACK/drivers/rl_cholqr_dense.hh new file mode 100644 index 000000000..c3153704d --- /dev/null +++ b/RandLAPACK/drivers/rl_cholqr_dense.hh @@ -0,0 +1,353 @@ +#pragma once + +#include "rl_util.hh" +#include "rl_blaspp.hh" +#include "rl_lapackpp.hh" +#include "rl_linops.hh" +#include "rl_exceptions.hh" +#include "../comps/rl_cholqr.hh" + +#include +#include +#include +#include + +namespace RandLAPACK { + +/// Dense (non-LinOp) entry points for the CholQR family. +/// +/// These drivers take a raw column-major buffer instead of a LinearOperator, for +/// callers that work through the regular BLAS API and do not want to build an +/// operator first. They run the SAME numerics as their LinOp counterparts in +/// `rl_cholqr_linops.hh` and `rl_scholqr3_linops.hh`: each wraps its input in a +/// `linops::DenseLinOp` and delegates to the shared `cholqr_iterate` engine, so +/// the pass count, the shift policy, and the adaptive-shift retry are inherited +/// rather than reimplemented. The only thing that differs is the interface. +/// +/// The three members of the family are distinguished exactly as in the engine: +/// CholQR_dense num_iters = 1, both shifts 0 (unshifted first attempt) +/// CholQR2_dense num_iters = 2, both shifts 0 +/// sCholQR3_dense num_iters = 3, iter 1 shifted by eps, iters 2-3 unshifted +/// +/// A is not modified. R is the (upper-triangular) output factor. When a Q buffer +/// is supplied, Q = A * R^{-1} is materialized after the factorization and is +/// excluded from the reported timing total, matching the LinOp drivers' test mode. + +namespace detail { + +/// Shared body for the three dense drivers. Wraps A in a DenseLinOp and runs +/// `num_iters` passes of the CholQR engine. +/// +/// @param[in] m Rows of A. Must be >= 0. +/// @param[in] n Columns of A. Must be >= 0 and <= m. +/// @param[in] A Column-major input buffer, not modified. +/// @param[in] lda Leading dimension of A, must be >= m. +/// @param[out] R Output factor, n by n, column-major. +/// @param[in] ldr Leading dimension of R, must be >= n. +/// @param[out] Q Optional. If non-null, receives Q = A * R^{-1} (m by n). +/// @param[in] ldq Leading dimension of Q. Must be >= m. Ignored when Q is null. +/// @param[out] applied_shifts Optional, length num_iters: per-pass absolute shift +/// the successful potrf carried (0 = unshifted). +/// @param[out] gram_traces Optional, length num_iters: per-pass Gram trace. +/// @param[out] q_mat_us Optional. When Q is materialized and timing is on, receives +/// the wall-clock (us) spent materializing Q, so the caller can +/// exclude it from the reported total exactly as the LinOp +/// drivers do. Set to 0 when Q is null or timing is off. +/// @return 0 on success, or the 1-based index of the pass whose Cholesky failed. +template +int cholqr_dense_body( + int64_t m, + int64_t n, + const T* A, + int64_t lda, + T* R, + int64_t ldr, + T* Q, + int64_t ldq, + int64_t block_size, + int num_iters, + T shift_iter1, + T shift_iter_rest, + int max_retries, + T shift_growth, + bool timing, + long* iter_times, + int* n_retries_total, + T* applied_shifts, + T* gram_traces, + long* q_mat_us +) { + randlapack_require(m >= 0) << "m=" << m << " must be >= 0"; + randlapack_require(n >= 0) << "n=" << n << " must be >= 0"; + randlapack_require(n <= m) << "n=" << n << " must be <= m=" << m << " (CholQR needs a tall or square input)"; + randlapack_require(lda >= m) << "lda=" << lda << " < m=" << m << " (column-major input)"; + randlapack_require(ldr >= n) << "ldr=" << ldr << " < n=" << n; + randlapack_require(!(A == nullptr && m > 0 && n > 0)) << "A buffer is null but m=" << m << " and n=" << n << " imply a nonempty matrix"; + randlapack_require(!(Q != nullptr && ldq < m)) << "ldq=" << ldq << " < m=" << m << " (column-major Q output)"; + + if (q_mat_us) *q_mat_us = 0; + + linops::DenseLinOp A_op(m, n, A, lda, Layout::ColMajor); + + int info = cholqr_iterate>( + A_op, R, ldr, block_size, num_iters, + shift_iter1, shift_iter_rest, + max_retries, shift_growth, timing, + iter_times, n_retries_total, applied_shifts, gram_traces); + if (info != 0) + return info; + + if (Q != nullptr) { + int64_t b_eff = (block_size > 0 && block_size < n) ? block_size : n; + if (timing && q_mat_us) { + std::chrono::steady_clock::time_point qt0 = std::chrono::steady_clock::now(); + RandLAPACK::materialize_Q_from_R(A_op, R, ldr, m, n, b_eff, Q, ldq); + std::chrono::steady_clock::time_point qt1 = std::chrono::steady_clock::now(); + *q_mat_us = std::chrono::duration_cast(qt1 - qt0).count(); + } else { + RandLAPACK::materialize_Q_from_R(A_op, R, ldr, m, n, b_eff, Q, ldq); + } + } + return 0; +} + +} // end namespace detail + + +/// Plain CholeskyQR on a dense column-major buffer. +/// One unpreconditioned pass, unshifted first attempt. +template +class CholQR_dense { + public: + + bool timing; + int64_t block_size; + + // Adaptive-shift safety net: the first attempt is unshifted, and only on + // potrf breakdown does the primitive seed the shift and grow it. A negative + // max_retries means unbounded, matching the LinOp drivers. + int max_retries; + T shift_growth; + int n_chol_retries = 0; ///< shift retries used on the last call (0 = clean) + /// Per-pass shift record from the last call: the absolute diagonal shift the + /// successful potrf carried (0 = unshifted) and that pass's Gram trace. + T chol_applied_shifts[1] = {T(0)}; + T chol_gram_traces[1] = {T(0)}; + + // 6 entries: alloc, fwd, adj, chol, rest, total + std::vector times; + long total_us() const { return times.empty() ? -1L : times.back(); } + + CholQR_dense( + bool time_subroutines + ) { + timing = time_subroutines; + block_size = kDefaultGramBlockSize; + max_retries = -1; + shift_growth = T(10); + } + + int call( + int64_t m, + int64_t n, + const T* A, + int64_t lda, + T* R, + int64_t ldr, + T* Q = nullptr, + int64_t ldq = 0 + ) { + std::chrono::steady_clock::time_point total_t_start, total_t_stop; + if (this->timing) total_t_start = std::chrono::steady_clock::now(); + + long it[5] = {0}; + long q_mat_us = 0; + int info = detail::cholqr_dense_body( + m, n, A, lda, R, ldr, Q, ldq, + this->block_size, /*num_iters=*/1, + /*shift_iter1=*/T(0), /*shift_iter_rest=*/T(0), + this->max_retries, this->shift_growth, this->timing, + this->timing ? it : nullptr, &this->n_chol_retries, + this->chol_applied_shifts, this->chol_gram_traces, + this->timing ? &q_mat_us : nullptr); + if (info != 0) return info; + + if (this->timing) { + total_t_stop = std::chrono::steady_clock::now(); + long total_dur = std::chrono::duration_cast(total_t_stop - total_t_start).count() - q_mat_us; + long fwd = it[0], adj = it[1], chol = it[3]; + long rest_dur = total_dur - (fwd + adj + chol); + this->times = {0L, fwd, adj, chol, rest_dur, total_dur}; + } + return 0; + } +}; + + +/// CholeskyQR2 on a dense column-major buffer. +/// Two passes, both starting unshifted; the retry rescues a non-PD Gram. +template +class CholQR2_dense { + public: + + bool timing; + int64_t block_size; + + int max_retries; + T shift_growth; + int n_chol_retries = 0; + /// Per-pass shift record from the last call (pass 1, pass 2): absolute shift + /// the successful potrf carried (0 = unshifted) and that pass's Gram trace. + T chol_applied_shifts[2] = {T(0), T(0)}; + T chol_gram_traces[2] = {T(0), T(0)}; + + // 11 entries: alloc, fwd1, adj1, chol1, upd1, fwd2, adj2, gemm2, chol2, upd2, total + // upd1 is 0 by convention (iter 1 has no R-update step). Matches + // CholQR2_linops's layout exactly (the always-zero iter-1 gemm slot is + // dropped, not carried as a 12th entry). + std::vector times; + long total_us() const { return times.empty() ? -1L : times.back(); } + + CholQR2_dense( + bool time_subroutines + ) { + timing = time_subroutines; + block_size = kDefaultGramBlockSize; + max_retries = -1; + shift_growth = T(10); + } + + int call( + int64_t m, + int64_t n, + const T* A, + int64_t lda, + T* R, + int64_t ldr, + T* Q = nullptr, + int64_t ldq = 0 + ) { + std::chrono::steady_clock::time_point total_t_start, total_t_stop; + if (this->timing) total_t_start = std::chrono::steady_clock::now(); + + long it[10] = {0}; + long q_mat_us = 0; + int info = detail::cholqr_dense_body( + m, n, A, lda, R, ldr, Q, ldq, + this->block_size, /*num_iters=*/2, + /*shift_iter1=*/T(0), /*shift_iter_rest=*/T(0), + this->max_retries, this->shift_growth, this->timing, + this->timing ? it : nullptr, &this->n_chol_retries, + this->chol_applied_shifts, this->chol_gram_traces, + this->timing ? &q_mat_us : nullptr); + if (info != 0) return info; + + if (this->timing) { + total_t_stop = std::chrono::steady_clock::now(); + long total_dur = std::chrono::duration_cast(total_t_stop - total_t_start).count() - q_mat_us; + // it = [fwd1,adj1,gemm1=0,chol1,upd1, fwd2,adj2,gemm2,chol2,upd2]. + // Drop the always-zero gemm1 slot (it[2]) to match CholQR2_linops. + this->times = {0L, it[0], it[1], it[3], it[4], + it[5], it[6], it[7], it[8], it[9], + total_dur}; + } + return 0; + } +}; + + +/// Shifted CholeskyQR3 on a dense column-major buffer. +/// Three passes; iter 1 carries the shift (default: the paper's 11*n*eps, or +/// eps via RANDLAPACK_SCHOLQR3_SHIFT=eps), iters 2 and 3 are unshifted +/// (Fukaya's prescription). This mirrors sCholQR3_linops exactly, including its +/// RANDLAPACK_SCHOLQR3_SHIFT env knob (shared via scholqr3_eps_shift() in +/// comps/rl_cholqr.hh). +template +class sCholQR3_dense { + public: + + bool timing; + int64_t block_size; + + T shift_factor_iter1; + T shift_factor_iter23; + + int max_retries; + T shift_growth; + int n_chol_retries = 0; + /// Per-pass shift record from the last call (passes 1-3): absolute shift the + /// successful potrf carried (0 = unshifted) and that pass's Gram trace. + T chol_applied_shifts[3] = {T(0), T(0), T(0)}; + T chol_gram_traces[3] = {T(0), T(0), T(0)}; + + // Timing breakdown (18 entries; matches sCholQR3_linops exactly): + // [0] alloc + // [1] fwd1 [2] adj1 [3] chol1 [4] upd1 + // [5] fwd2 [6] adj2 [7] gemm2 [8] chol2 [9] upd2 + // [10] fwd3 [11] adj3 [12] gemm3 [13] chol3 [14] upd3 + // [15] q_mat [16] rest [17] total + std::vector times; + long total_us() const { return times.empty() ? -1L : times.back(); } + + sCholQR3_dense( + bool time_subroutines + ) { + timing = time_subroutines; + block_size = kDefaultGramBlockSize; + shift_factor_iter1 = T(-1); // < 0: resolve default (11*n*eps, or eps via env) at call time + shift_factor_iter23 = T(0); + max_retries = -1; + shift_growth = T(10); + } + + int call( + int64_t m, + int64_t n, + const T* A, + int64_t lda, + T* R, + int64_t ldr, + T* Q = nullptr, + int64_t ldq = 0 + ) { + std::chrono::steady_clock::time_point total_t_start, total_t_stop; + if (this->timing) total_t_start = std::chrono::steady_clock::now(); + + // First-pass shift defaults to the paper's s = 11*eps*n*trace(G) + // (FukayaEtAl2020, c = 11); RANDLAPACK_SCHOLQR3_SHIFT=eps selects the + // legacy eps*trace(G). A caller-set shift_factor_iter1 >= 0 wins. + // Same knob resolution as sCholQR3_linops. + const T eps_T = std::numeric_limits::epsilon(); + const T sf1 = (this->shift_factor_iter1 >= T(0)) + ? this->shift_factor_iter1 + : (scholqr3_eps_shift() ? eps_T : T(11) * T(n) * eps_T); + + long it[15] = {0}; + long q_mat_us = 0; + int info = detail::cholqr_dense_body( + m, n, A, lda, R, ldr, Q, ldq, + this->block_size, /*num_iters=*/3, + sf1, this->shift_factor_iter23, + this->max_retries, this->shift_growth, this->timing, + this->timing ? it : nullptr, &this->n_chol_retries, + this->chol_applied_shifts, this->chol_gram_traces, + this->timing ? &q_mat_us : nullptr); + if (info != 0) return info; // 1/2/3 = the pass that failed + + if (this->timing) { + total_t_stop = std::chrono::steady_clock::now(); + long total_dur = std::chrono::duration_cast(total_t_stop - total_t_start).count() - q_mat_us; + long iters_sum = 0; + for (int i = 0; i < 15; ++i) iters_sum += it[i]; + long rest_dur = total_dur - iters_sum; + this->times = {0L, + it[0], it[1], it[3], it[4], // fwd1, adj1, chol1, upd1 + it[5], it[6], it[7], it[8], it[9], // fwd2, adj2, gemm2, chol2, upd2 + it[10], it[11], it[12], it[13], it[14], // fwd3, adj3, gemm3, chol3, upd3 + q_mat_us, rest_dur, total_dur}; + } + return 0; + } +}; + +} // end namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_cholqr_linops.hh b/RandLAPACK/drivers/rl_cholqr_linops.hh index 5dc7a3494..62d5ba5cc 100644 --- a/RandLAPACK/drivers/rl_cholqr_linops.hh +++ b/RandLAPACK/drivers/rl_cholqr_linops.hh @@ -4,6 +4,7 @@ #include "rl_blaspp.hh" #include "rl_lapackpp.hh" #include "rl_linops.hh" +#include "../comps/rl_cholqr.hh" #include #include @@ -16,20 +17,12 @@ namespace RandLAPACK { /// Cholesky QR factorization for abstract linear operators. /// -/// Computes A = QR where A is any type satisfying the LinearOperator concept -/// (dense, sparse, composite, etc.). Unlike standard Cholesky QR (rl_cqrrt.hh), -/// which requires A to be a dense matrix that can be modified in place, this -/// class works through the operator interface: it only calls A(NoTrans, ...) -/// and A(Trans, ...) to form the Gram matrix G = A^T A, then factors G = R^T R -/// via Cholesky. +/// Thin wrapper around `cholqr_primitive` (Algorithm 1 from the collaborator's spec): +/// G = A^T A, G = R^T R via Cholesky. /// -/// This is useful when A is too large to store densely, is only available as a -/// matrix-vector product (e.g., the product of two factors, or a sparse matrix), -/// or when the caller wants a uniform interface across operator types. -/// -/// The Q factor is not computed by default (Q-less factorization). When -/// test_mode is enabled, Q = A * R^{-1} is materialized explicitly for -/// verification. +/// A may be any type satisfying the LinearOperator concept (dense, sparse, composite, +/// etc.). Q is not computed by default; when test_mode is enabled, Q = A * R^{-1} +/// is materialized for verification. /// template class CholQR_linops { @@ -37,7 +30,6 @@ class CholQR_linops { bool timing; bool test_mode; - T eps; // Q-factor for test mode (only allocated if test_mode = true) T* Q; @@ -45,48 +37,45 @@ class CholQR_linops { int64_t Q_cols; // 6 entries: alloc, fwd, adj, chol, rest, total - // fwd = LinOp NoTrans: A * I (accumulated over blocks) - // adj = LinOp Trans: A^T * buf (accumulated over blocks) std::vector times; - - // Column-block size for the materialize + Gram computation. - // - // When block_size > 0, the two expensive operations: - // (1) A_temp = A * I (m × n) - materialize the operator - // (2) R = A^T * A_temp (n × n) - Gram matrix - // are fused into a column-block loop that processes b columns at a time: - // for each column block j of width b: - // buf (m × b) = A * I[:, j*b : (j+1)*b] - // R[:, j*b : (j+1)*b] = A^T * buf - // - // This reduces peak memory from O(m*n) to O(m*b), which is significant - // when m is large and n is moderate. The result is mathematically - // identical — each column block of R is: - // R[:, j_block] = A^T * (A * I[:, j_block]) - // which equals the corresponding columns of A^T * A. - // - // When block_size <= 0 or block_size >= n, the full m × n buffer is - // allocated and the original (non-blocked) path is used. - // - // When test_mode is enabled and blocking is active, the Q-factor - // computation (which needs the full m × n A_temp) is handled by - // recomputing A_temp = A * I after the Gram loop. This - // recomputation is outside the timing region, so it does not - // affect benchmark results. + /// Total measured wall-clock (microseconds) of the last call(), or -1 if timing + /// was off. Every driver in this family packs the total as the LAST times[] entry, + /// but the entry COUNT differs per driver (6 / 11 / 15 / 18). Callers used to hard- + /// code that index (times[5], times[10], times[14], times[17]), so adding or + /// removing one slot silently wrote the wrong number into every CSV with no compile + /// error. Read the total through here instead. + long total_us() const { return times.empty() ? -1L : times.back(); } + + // Column-block size for Gram and Q materialization. <=0 or >=n means no blocking. int64_t block_size; + // Adaptive-shift safety net: the first attempt is always unshifted + // (shift_factor is hard-wired to 0 in call()); only if potrf breaks + // down does the primitive seed the shift at eps*trace(G) and grow it + // x shift_growth. max_retries < 0 = unbounded (no ceiling), retry until PD. + int max_retries; + T shift_growth; + int n_chol_retries = 0; ///< shift retries used on the last call (0 = clean) + /// Per-pass shift record from the last call: the absolute diagonal shift the + /// successful potrf carried (0 = unshifted) and that pass's Gram trace. A + /// nonzero shift means R factors G + s I, not G (preconditioner semantics). + T chol_applied_shifts[1] = {T(0)}; + T chol_gram_traces[1] = {T(0)}; + CholQR_linops( bool time_subroutines, T ep, bool enable_test_mode = false ) { timing = time_subroutines; - eps = ep; - block_size = 0; + (void)ep; // kept in the signature for call-site compatibility; unused + block_size = kDefaultGramBlockSize; test_mode = enable_test_mode; Q = nullptr; Q_rows = 0; Q_cols = 0; + max_retries = -1; // unbounded retries (no ceiling) + shift_growth = T(10); } ~CholQR_linops() { @@ -95,229 +84,199 @@ class CholQR_linops { } } - /// Computes an R-factor of the unpivoted QR factorization using unpreconditioned Cholesky QR: - /// A = QR, - /// where Q and R are of size m-by-n and n-by-n. - /// - /// This is the baseline unpreconditioned version for comparison with CQRRT. - /// Algorithm: - /// 1. Compute Gram matrix: G = A^T * A - /// 2. Compute Cholesky factorization: G = R^T * R - /// 3. (Optional) Compute Q = A * R^{-1} - /// - /// @note This algorithm expects A to be full-rank (rank = n). Rank-deficient inputs may result - /// in loss of orthogonality in the Q-factor (when test_mode=true) and numerical instability - /// in the R-factor. - /// - /// @param[in] A - /// The m-by-n linear operator (m and n read from A.n_rows, A.n_cols). - /// - /// @param[out] R - /// Pre-allocated n-by-n buffer. On exit, stores the upper-triangular - /// R factor. Zero entries are not compressed. - /// - /// @param[in] ldr - /// Leading dimension of R. - /// - /// @return = 0: successful exit template int call( GLO& A, T* R, int64_t ldr ) { - ///--------------------TIMING VARS--------------------/ - steady_clock::time_point alloc_t_start; - steady_clock::time_point alloc_t_stop; - steady_clock::time_point potrf_t_start; - steady_clock::time_point potrf_t_stop; - steady_clock::time_point total_t_start; - steady_clock::time_point total_t_stop; - steady_clock::time_point t_start, t_stop; - long alloc_t_dur = 0; - long fwd_t_dur = 0; - long adj_t_dur = 0; - long potrf_t_dur = 0; - long total_t_dur = 0; - long q_t_dur = 0; - steady_clock::time_point q_t_start; - steady_clock::time_point q_t_stop; - - if(this->timing) - total_t_start = steady_clock::now(); + steady_clock::time_point t0, t1, total_t_start, total_t_stop; + long q_dur = 0; + + if (this->timing) total_t_start = steady_clock::now(); int64_t m = A.n_rows; int64_t n = A.n_cols; - - // Compute Gram matrix: R = A^T * A - // We cannot use syrk since A may be a sparse operator - // Instead, compute R = A^T * A via two matvec operations: - // 1. Create identity matrix I (n x n) - // 2. Compute A_temp = A * I (materializes A as m x n dense) - // 3. Compute R = A^T * A_temp - - if(this->timing) - alloc_t_start = steady_clock::now(); - - // Create identity matrix (needs zero-init for off-diagonal elements) - T* I_mat = new T[n * n](); - RandLAPACK::util::eye(n, n, I_mat); - - // Gram computation: R = A^T * A - - // Determine effective block width. - // block_size <= 0 or >= n means "no blocking" (full width). int64_t b_eff = (this->block_size > 0 && this->block_size < n) ? this->block_size : n; - // A_temp buffer: - // Full path: m × n (kept alive for Q-factor in test_mode) - // Block path: m × b_eff (temporary, freed after Gram loop; - // if test_mode, a full m × n buffer is - // allocated later for Q computation) - // No zero-init needed: first use is with beta=0.0 which overwrites all elements. - T* A_temp = new T[m * b_eff]; - - if(this->timing) { - alloc_t_stop = steady_clock::now(); + // Plain CholQR = one unpreconditioned cholqr_iterate pass, unshifted-first. + long it[5] = {0}; + int info = cholqr_iterate( + A, R, ldr, this->block_size, /*num_iters=*/1, + /*shift_iter1=*/T(0), /*shift_iter_rest=*/T(0), + this->max_retries, this->shift_growth, this->timing, + this->timing ? it : nullptr, &this->n_chol_retries, + this->chol_applied_shifts, this->chol_gram_traces); + // 1 = the (only) pass failed; the cause (retry exhaustion, singular + // preconditioner, non-finite shift, invalid input) is on stderr. + if (info != 0) return info; + + // Test mode: materialize Q = A * R^{-1} (outside the timing region). + if (this->test_mode) { + if (this->timing) t0 = steady_clock::now(); + T* Q_buf = new T[m * n]; + RandLAPACK::materialize_Q_from_R(A, R, ldr, m, n, b_eff, Q_buf, m); + this->Q_rows = m; + this->Q_cols = n; + // The class owns Q (the destructor frees it), so release any buffer from a + // previous call() before taking ownership of this one. + delete[] this->Q; + this->Q = Q_buf; + if (this->timing) { t1 = steady_clock::now(); q_dur = duration_cast(t1 - t0).count(); } } - if (b_eff == n) { - // --- Full materialization path (original) --- - - // Step 1: Materialize A by computing A_temp = A * I - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, m, n, n, (T)1.0, I_mat, n, (T)0.0, A_temp, m); - if(this->timing) { t_stop = steady_clock::now(); fwd_t_dur = duration_cast(t_stop - t_start).count(); } - - // Step 2: Compute R = A^T * A_temp (using the linear operator's transpose) - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, n, n, m, (T)1.0, A_temp, m, (T)0.0, R, ldr); - if(this->timing) { t_stop = steady_clock::now(); adj_t_dur = duration_cast(t_stop - t_start).count(); } - } else { - // --- Column-block processing path (memory-efficient) --- - // - // Process n columns in blocks of width b_eff. - // The last block may be narrower if n is not divisible by b_eff. - // - // For each block starting at column j with width b_j: - // (1) buf (m × b_j) = A * I[:, j : j+b_j] - // I is n × n in ColMajor with ld = n. - // Column j starts at I + j * n. - // We multiply A (m × n) by this n × b_j slice. - // - // (2) R[:, j : j+b_j] (n × b_j) = A^T * buf - // A^T is n × m, buf is m × b_j, result is n × b_j. - // R is n × n with ld = ldr. - // Column j starts at R + j * ldr. - // - // Total FLOPs are the same as the full path; only memory differs. - - long fwd_accum = 0, adj_accum = 0; - for (int64_t j = 0; j < n; j += b_eff) { - int64_t b_j = std::min(b_eff, n - j); - - // (1) buf = A * I[:, j : j+b_j] - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, b_j, n, (T)1.0, I_mat + j * n, n, (T)0.0, A_temp, m); - if(this->timing) { t_stop = steady_clock::now(); fwd_accum += duration_cast(t_stop - t_start).count(); } - - // (2) R[:, j : j+b_j] = A^T * buf - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, - n, b_j, m, (T)1.0, A_temp, m, (T)0.0, R + j * ldr, ldr); - if(this->timing) { t_stop = steady_clock::now(); adj_accum += duration_cast(t_stop - t_start).count(); } - } - if(this->timing) { - fwd_t_dur = fwd_accum; - adj_t_dur = adj_accum; - } + if (this->timing) { + total_t_stop = steady_clock::now(); + long total_dur = duration_cast(total_t_stop - total_t_start).count() - q_dur; + long fwd = it[0], adj = it[1], chol = it[3]; + long rest_dur = total_dur - (fwd + adj + chol); + this->times = {0L, fwd, adj, chol, rest_dur, total_dur}; // [alloc, fwd, adj, chol, rest, total] } - // Zero out the lower triangle before Cholesky (potrf only uses upper triangle) - if (n > 1) { - lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &R[1], ldr); - } + return 0; + } +}; - if(this->timing) { - potrf_t_start = steady_clock::now(); - } - // Compute Cholesky factorization: G = R^T * R - // On exit, the upper triangle of R contains the R-factor - if (lapack::potrf(Uplo::Upper, n, R, ldr)) { - delete[] I_mat; - delete[] A_temp; - return 1; - } +/// CholQR2 for abstract linear operators. +/// +/// Two passes of CholQR via the shared primitives: +/// iter 1: cholqr_primitive(A, shift_factor, max_retries) -> R_1 +/// iter 2: cholqr_primitive(A, P=R_1, TRSM_IDENTITY, ..., max_retries) -> R +/// +/// Both passes start UNSHIFTED (shift_factor = 0); only on potrf breakdown does +/// the primitive seed the shift at eps * trace(G) (= ||A||_F^2 for iter 1, ~ n for +/// the iter-2 preconditioned Gram) and grow it x shift_growth, retrying unboundedly +/// (max_retries < 0) until the Gram is PD. Starting unshifted +/// avoids biasing R_1 on well-conditioned inputs: an always-on eps shift was found +/// to leave CholQR2 *less* orthogonal than a single unshifted CholQR pass; the retry +/// still rescues Gram matrices driven non-PD by rounding. +/// +/// Status codes from call(): the 1-based pass whose factorization failed, 0 on +/// success. The failure can have any cause cholqr_primitive reports (retry +/// exhaustion, a singular preconditioner, a non-finite shift, or invalid +/// input); see stderr for which one fired. +/// 1 pass 1 (unpreconditioned CholQR) failed +/// 2 pass 2 (preconditioned on R_1) failed +/// +template +class CholQR2_linops { + public: + bool timing; + bool test_mode; - if(this->timing) - potrf_t_stop = steady_clock::now(); - - // Compute Q-factor if test mode is enabled (NOT included in cholqr timing) - if(this->test_mode) { - if(this->timing) - q_t_start = steady_clock::now(); - - if (b_eff < n) { - // Column-block Gram was used: A_temp is only m × b_eff, - // too small for Q. Recompute A_temp = A * I in - // full. This is outside the timing region, so the extra - // operator application does not affect benchmark results. - delete[] A_temp; - // No zero-init: beta=0.0 overwrites all elements - A_temp = new T[m * n]; - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, n, n, (T)1.0, I_mat, n, (T)0.0, A_temp, m); - } + // Q-factor for test mode (only allocated if test_mode = true) + T* Q; + int64_t Q_rows; + int64_t Q_cols; - this->Q_rows = m; - this->Q_cols = n; - this->Q = A_temp; // Take ownership of A_temp buffer + // 11 entries: alloc, fwd1, adj1, chol1, upd1, fwd2, adj2, gemm2, chol2, upd2, total + // upd1 is 0 by convention (iter 1 has no R-update step). + std::vector times; + /// Total measured wall-clock (microseconds) of the last call(), or -1 if timing + /// was off. Every driver in this family packs the total as the LAST times[] entry, + /// but the entry COUNT differs per driver (6 / 11 / 15 / 18). Callers used to hard- + /// code that index (times[5], times[10], times[14], times[17]), so adding or + /// removing one slot silently wrote the wrong number into every CSV with no compile + /// error. Read the total through here instead. + long total_us() const { return times.empty() ? -1L : times.back(); } - // Solve Q * R = A_temp for Q - // Q = A * R^{-1} - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, m, n, (T)1.0, R, ldr, this->Q, m); + int64_t block_size; - if(this->timing) - q_t_stop = steady_clock::now(); - } + // Adaptive-shift policy (see cholqr_primitive). + // shift_factor_iter1 is multiplied by trace(G_1) = ||A||_F^2 on the first + // attempt. shift_factor_iter2 is multiplied by trace(G_2) (~ n when the + // preconditioner is well-formed). max_retries bounds the geometric retry + // loop; final shift is shift_factor * shift_growth^k on the kth retry. + T shift_factor_iter1; + T shift_factor_iter2; + int max_retries; + T shift_growth; + int n_chol_retries = 0; ///< shift retries used on the last call (0 = clean) + /// Per-pass shift record from the last call (pass 1, pass 2): absolute shift + /// the successful potrf carried (0 = unshifted) and that pass's Gram trace. + T chol_applied_shifts[2] = {T(0), T(0)}; + T chol_gram_traces[2] = {T(0), T(0)}; + + CholQR2_linops( + bool time_subroutines, + T ep, + bool enable_test_mode = false + ) { + timing = time_subroutines; + (void)ep; // kept in the signature for call-site compatibility; unused + block_size = kDefaultGramBlockSize; + test_mode = enable_test_mode; + Q = nullptr; + Q_rows = 0; + Q_cols = 0; + shift_factor_iter1 = T(0); // unshifted first attempt (shift only on breakdown) + shift_factor_iter2 = T(0); + max_retries = -1; // unbounded retries (no ceiling) + shift_growth = T(10); + } - if(this->timing) { - // Stop timing BEFORE cleanup operations to exclude deallocation costs - total_t_stop = steady_clock::now(); + ~CholQR2_linops() { + if (Q != nullptr) delete[] Q; + } - alloc_t_dur = duration_cast(alloc_t_stop - alloc_t_start).count(); - // fwd_t_dur and adj_t_dur already set (in both full and blocked paths) - potrf_t_dur = duration_cast(potrf_t_stop - potrf_t_start).count(); - total_t_dur = duration_cast(total_t_stop - total_t_start).count(); + template + int call( + GLO& A, + T* R, + int64_t ldr + ) { + steady_clock::time_point t0, t1, total_t_start, total_t_stop; + long q_mat_dur = 0; - // Subtract Q-factor computation time if in test mode - if(this->test_mode) { - q_t_dur = duration_cast(q_t_stop - q_t_start).count(); - total_t_dur -= q_t_dur; - } + if (this->timing) total_t_start = steady_clock::now(); - long rest_t_dur = total_t_dur - (alloc_t_dur + fwd_t_dur + adj_t_dur + potrf_t_dur); + int64_t m = A.n_rows; + int64_t n = A.n_cols; + int64_t b_eff = (this->block_size > 0 && this->block_size < n) + ? this->block_size : n; - // Fill the data vector: [alloc, fwd, adj, chol, rest, total] - this->times = {alloc_t_dur, fwd_t_dur, adj_t_dur, potrf_t_dur, rest_t_dur, total_t_dur}; + // CholQR2 = two cholqr_iterate passes (iter 1 unpreconditioned, iter 2 + // preconditioned on R_1). Both shifts default to 0 (unshifted-first). + long it[10] = {0}; + int info = cholqr_iterate( + A, R, ldr, this->block_size, /*num_iters=*/2, + this->shift_factor_iter1, this->shift_factor_iter2, + this->max_retries, this->shift_growth, this->timing, + this->timing ? it : nullptr, &this->n_chol_retries, + this->chol_applied_shifts, this->chol_gram_traces); + // 1 or 2 = the 1-based pass that failed (retry exhaustion, singular + // preconditioner, non-finite shift, or invalid input; see stderr). + if (info != 0) return info; + + // ---- Test mode: materialize Q = A * R^{-1} via blocked linop calls ---- + if (this->test_mode) { + if (this->timing) t0 = steady_clock::now(); + T* Q_buf = new T[m * n]; + RandLAPACK::materialize_Q_from_R(A, R, ldr, m, n, b_eff, Q_buf, m); + this->Q_rows = m; + this->Q_cols = n; + // The class owns Q (the destructor frees it), so release any buffer from a + // previous call() before taking ownership of this one. + delete[] this->Q; + this->Q = Q_buf; + if (this->timing) { t1 = steady_clock::now(); q_mat_dur = duration_cast(t1 - t0).count(); } } - // Cleanup - now outside the timing region to avoid timing artifacts - delete[] I_mat; - - // Only delete A_temp if not in test mode (otherwise Q owns it). - // When test_mode + blocking: the small block buffer was freed - // and replaced with a full m × n buffer in the Q section above. - if(!this->test_mode) { - delete[] A_temp; + if (this->timing) { + total_t_stop = steady_clock::now(); + long total_dur = duration_cast(total_t_stop - total_t_start).count() - q_mat_dur; + // it = [fwd1,adj1,gemm1=0,chol1,upd1=0, fwd2,adj2,gemm2,chol2,upd2]. + this->times = {0L, // alloc (now inside cholqr_iterate) + it[0], it[1], it[3], it[4], // fwd1, adj1, chol1, upd1 + it[5], it[6], it[7], it[8], it[9], // fwd2, adj2, gemm2, chol2, upd2 + total_dur}; } return 0; } }; + } // end namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_cqrrt.hh b/RandLAPACK/drivers/rl_cqrrt.hh index d1dbc55f1..2731f006c 100644 --- a/RandLAPACK/drivers/rl_cqrrt.hh +++ b/RandLAPACK/drivers/rl_cqrrt.hh @@ -6,17 +6,36 @@ #include "rl_lapackpp.hh" #include "rl_hqrrp.hh" #include "rl_bqrrp.hh" +#include "rl_linops.hh" +#include "../comps/rl_cholqr.hh" #include #include +#include #include #include #include +#include +#include using namespace std::chrono; namespace RandLAPACK { +/// Backwards-compatible alias for the precond-method enum, which now lives in +/// comps/rl_cholqr.hh as PCholQRPrecondMethod (shared across CholQR/sCholQR3/CQRRT). +using CQRRTLinopPrecond = PCholQRPrecondMethod; + + +// ============================================================================ +// CQRRT: dense Q-less Cholesky QR with sketch preconditioning. +// ============================================================================ +// +// Operates on a dense column-major A (m × n). Modifies A in place via TRSM +// (A := A * R_sk^{-1}) so the Q factor is materialized implicitly inside A. +// This is faster than the linop variant when A is already dense in memory. +// +// Reference: arXiv:2111.11148. template class CQRRTalg { public: @@ -47,47 +66,27 @@ class CQRRT : public CQRRTalg { eps = ep; orthogonalization = false; compute_Q = true; - nnz = 2; + nnz = 4; // SASO nonzeros per column; paper uses 4 or 8 (2 causes sporadic spikes) + max_retries = -1; // unbounded retries (no ceiling), as CQRRT_linops + shift_growth = T(10); } /// Computes an unpivoted QR factorization of the form: /// A= QR, /// where Q and R are of size m-by-n and n-by-n. - /// Detailed description of this algorithm may be found in https://arxiv.org/pdf/2111.11148. /// /// @note This algorithm expects A to be full-rank (rank = n). Rank-deficient inputs may result /// in loss of orthogonality in the Q-factor and numerical instability in the R-factor. /// - /// @param[in] m - /// The number of rows in the matrix A. - /// - /// @param[in] n - /// The number of columns in the matrix A. - /// - /// @param[in] A - /// The m-by-n matrix A, stored in a column-major format. - /// - /// @param[in] d - /// Embedding dimension of a sketch, m >= d >= n. - /// - /// @param[in] R - /// Represents the upper-triangular R factor of QR factorization. - /// On entry, is empty and may not have any space allocated for it. - /// - /// @param[in] state - /// RNG state parameter, required for sketching operator generation. - /// - /// @param[out] A - /// Overwritten by an m-by-n orthogonal Q factor. - /// Matrix is stored explicitly. - /// - /// @param[out] R - /// Stores n-by-n matrix with upper-triangular R factor. - /// Zero entries are not compressed. - /// - /// @return = 0: successful exit - /// - + /// @return 0 on success; 1 if the preconditioner's diagonal is singular or the + /// preconditioned-Gram Cholesky (potrf) failed on every shift-retry + /// attempt (see stderr for which); -2 if shift_growth is invalid + /// (<= 1, which defeats the retry's geometric-growth termination + /// argument), a caller-configuration bug, checked before any work, + /// same sentinel cholqr_primitive uses for its own invalid-shift + /// inputs. Dense CQRRT has no caller-exposed shift_factor to validate + /// separately: the first attempt always starts unshifted, same as + /// CQRRT_linops, so shift_growth is the only invalid-shift input here. int call( int64_t m, int64_t n, @@ -103,19 +102,35 @@ class CQRRT : public CQRRTalg { bool timing; T eps; - // 10 entries: saso, qr, trtri(=0), precond, gram, trmm_gram(=0), potrf, finalize, rest, total - // Matches CQRRT_linops timing indices for direct comparison. + // 10 entries: saso, qr, trtri(=0), precond, gram, trmm_gram(=0), potrf, finalize, rest, total. + // NOT index-compatible with CQRRT_linops's 11-entry layout (that one starts + // with an alloc slot); a plotter must dispatch on the layout, not assume the + // indices line up. std::vector times; + /// Total measured wall-clock (microseconds) of the last call(), or -1 if timing + /// was off. Every driver in this family packs the total as the LAST times[] entry, + /// but the entry COUNT differs per driver (6 / 11 / 15 / 18). Callers used to hard- + /// code that index (times[5], times[10], times[14], times[17]), so adding or + /// removing one slot silently wrote the wrong number into every CSV with no compile + /// error. Read the total through here instead. + long total_us() const { return times.empty() ? -1L : times.back(); } - // tuning SASOS int64_t nnz; - - // Mode of operation that allows to use CQRRT for orthogonalization of the input matrix. bool orthogonalization; - - // If false, skip the Q-factor computation (R-only mode). - // When false, A is NOT overwritten with Q on output. - bool compute_Q; + bool compute_Q; // skip Q materialization when false (R-only mode) + + // Adaptive-shift safety net on the preconditioned Gram's Cholesky, matching + // CQRRT_linops: the first attempt is always unshifted; only on potrf + // breakdown does the retry seed the shift at eps*trace(G) and grow it + // x shift_growth. max_retries < 0 = unbounded (retry until PD). The clean + // (unshifted, first-attempt-succeeds) path is unaffected. + int max_retries; + T shift_growth; + int n_chol_retries = 0; ///< shift retries used on the last call (0 = clean) + /// Shift record from the last call's preconditioned-Gram Cholesky: absolute + /// shift the successful potrf carried (0 = unshifted) and the Gram's trace. + T chol_applied_shifts[1] = {T(0)}; + T chol_gram_traces[1] = {T(0)}; }; // ----------------------------------------------------------------------------- @@ -135,11 +150,34 @@ int CQRRT::call( // CQRRT is called through a binding layer (e.g. MEX/MATLAB). randlapack_require(m >= 0) << "m=" << m << " must be >= 0"; randlapack_require(n >= 0) << "n=" << n << " must be >= 0"; + randlapack_require(m >= n) << "CQRRT: operator must be tall (m=" << m << " < n=" << n << ")"; randlapack_require(lda >= m) << "lda=" << lda << " < m=" << m << " (lda must be >= m for ColMajor)"; randlapack_require(ldr >= n) << "ldr=" << ldr << " < n=" << n << " (ldr must be >= n)"; randlapack_require(d_factor >= (T)1.0) << "d_factor=" << d_factor << " must be >= 1.0"; randlapack_require(!(A == nullptr && m > 0 && n > 0)) << "A buffer is null but m=" << m << " and n=" << n << " imply a nonempty matrix"; - randlapack_require(!(R == nullptr && n > 0)) << "R buffer is null but n=" << n << " > 0"; + randlapack_require(R != nullptr) << "CQRRT: R buffer is null"; + + // Reset the shift-record out-params up front so a stale value from a + // previous call never survives an early-return failure path. + this->n_chol_retries = 0; + this->chol_applied_shifts[0] = T(0); + this->chol_gram_traces[0] = T(0); + + // Fail-fast shift-config validation, same argument as cholqr_primitive's + // shift check (rl_cholqr.hh): shift_growth <= 1 defeats the retry loop's + // geometric-growth termination argument (it would spin at a constant or + // shrinking shift instead of escalating toward diagonal dominance). No + // separate shift_factor check is needed here: unlike cholqr_primitive, + // dense CQRRT does not expose a caller-supplied starting shift_factor + // (the first attempt is always unshifted, matching CQRRT_linops), so + // there is nothing else to validate before the retry loop below. + if (this->shift_growth <= T(1)) { + std::fprintf(stderr, + "[CQRRT] FAIL: shift_growth (%g) must be > 1 (<= 1 defeats the " + "geometric-growth retry termination argument)\n", + (double)this->shift_growth); + return -2; + } ///--------------------TIMING VARS--------------------/ steady_clock::time_point saso_t_start, saso_t_stop; @@ -150,131 +188,125 @@ int CQRRT::call( steady_clock::time_point q_t_start, q_t_stop; steady_clock::time_point finalize_t_start, finalize_t_stop; steady_clock::time_point total_t_start, total_t_stop; - long saso_t_dur = 0; - long qr_t_dur = 0; - long precond_t_dur = 0; - long gram_t_dur = 0; - long potrf_t_dur = 0; - long q_t_dur = 0; - long finalize_t_dur = 0; - long total_t_dur = 0; - - if(this -> timing) - total_t_start = steady_clock::now(); + long saso_t_dur = 0, qr_t_dur = 0, precond_t_dur = 0, gram_t_dur = 0; + long potrf_t_dur = 0, q_t_dur = 0, finalize_t_dur = 0, total_t_dur = 0; - int64_t d = d_factor * n; + if(this -> timing) total_t_start = steady_clock::now(); + int64_t d = (int64_t) (d_factor * (T) n); + if (d < n) d = n; // same clamp as CQRRT_linops: truncation must not undershoot n T* A_hat = new T[d * n](); T* tau = new T[n](); - if(this -> timing) - saso_t_start = steady_clock::now(); - - /// Generating a SASO + // Sketch + small QR + if(this -> timing) saso_t_start = steady_clock::now(); RandBLAS::SparseDist DS(d, m, this->nnz); RandBLAS::SparseSkOp S(DS, state); state = S.next_state; + RandBLAS::sketch_general(Layout::ColMajor, Op::NoTrans, Op::NoTrans, + d, n, m, (T)1.0, S, 0, 0, A, lda, (T)0.0, A_hat, d); + if(this -> timing) { saso_t_stop = steady_clock::now(); qr_t_start = steady_clock::now(); } - /// Applying a SASO - RandBLAS::sketch_general( - Layout::ColMajor, Op::NoTrans, Op::NoTrans, - d, n, m, (T) 1.0, S, 0, 0, A, lda, (T) 0.0, A_hat, d - ); - - if(this -> timing) { - saso_t_stop = steady_clock::now(); - qr_t_start = steady_clock::now(); - } - - /// Performing QR on a sketch lapack::geqrf(d, n, A_hat, d, tau); + if(this -> timing) qr_t_stop = steady_clock::now(); - if(this -> timing) - qr_t_stop = steady_clock::now(); - - /// Extracting a k by k R representation - T* R_sk = R; + T* R_sk = R; lapack::lacpy(MatrixType::Upper, n, n, A_hat, d, R_sk, ldr); - - if(this -> timing) - precond_t_start = steady_clock::now(); - - // Precondition: A := A * R_sk^{-1} + // The caller's R buffer is otherwise uninitialized; the trmm below reads all + // of R_sk (not just its upper triangle), so a garbage strict lower would + // contaminate the output R's upper triangle (mirrors cholqr_primitive's own + // lacpy+laset pattern at rl_cholqr.hh). + if (n > 1) + lapack::laset(MatrixType::Lower, n - 1, n - 1, T(0), T(0), R_sk + 1, ldr); + + if(this -> timing) precond_t_start = steady_clock::now(); if (!RandLAPACK::util::diag_is_nonzero(n, R_sk, ldr)) { - delete[] A_hat; - delete[] tau; - return 1; - } - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, Diag::NonUnit, m, n, 1.0, R_sk, ldr, A, lda); - - if(this -> timing) { - precond_t_stop = steady_clock::now(); - gram_t_start = steady_clock::now(); + delete[] A_hat; delete[] tau; return 1; } + blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, Diag::NonUnit, + m, n, 1.0, R_sk, ldr, A, lda); + if(this -> timing) { precond_t_stop = steady_clock::now(); gram_t_start = steady_clock::now(); } - // Gram matrix: G = A^T * A (SYRK, upper triangle only) blas::syrk(Layout::ColMajor, Uplo::Upper, Op::Trans, n, m, 1.0, A, lda, 0.0, R_sk, ldr); - - if(this -> timing) { - gram_t_stop = steady_clock::now(); - potrf_t_start = steady_clock::now(); + if(this -> timing) { gram_t_stop = steady_clock::now(); potrf_t_start = steady_clock::now(); } + + // Adaptive-shift retry (same policy as cholqr_primitive's Step 3, parity with + // CQRRT_linops): the clean path (unshifted potrf succeeds first try) is still + // bit-identical to before, but is NOT free of extra work. With the default + // max_retries = -1 (unbounded), gram_backup below is allocated and filled via + // lacpy on every call, clean or not; that O(n^2) alloc+copy is timed into the + // potrf slot even when no retry ever happens. Only max_retries == 0 skips it. + T* gram_backup = (this->max_retries != 0) ? new T[n * n] : nullptr; + T trace_G = 0; + for (int64_t i = 0; i < n; ++i) trace_G += R_sk[i * (ldr + 1)]; + this->chol_gram_traces[0] = trace_G; + if (gram_backup) lapack::lacpy(MatrixType::Upper, n, n, R_sk, ldr, gram_backup, n); + + int potrf_info = 0; + int attempt = 0; + T current_shift_factor = T(0); + T last_shift = T(0); + constexpr int kUnboundedRetryCeiling = 128; + for (; (this->max_retries < 0) ? (attempt < kUnboundedRetryCeiling) : (attempt <= this->max_retries); ++attempt) { + if (attempt > 0) { + lapack::lacpy(MatrixType::Upper, n, n, gram_backup, n, R_sk, ldr); + current_shift_factor = (current_shift_factor > T(0)) + ? current_shift_factor * this->shift_growth + : std::numeric_limits::epsilon(); + if (!std::isfinite(current_shift_factor) || !std::isfinite(trace_G)) { + potrf_info = -1; + break; + } + } + if (current_shift_factor > T(0)) { + T shift = current_shift_factor * trace_G; + if (!std::isfinite(shift)) { + potrf_info = -1; + break; + } + for (int64_t i = 0; i < n; ++i) R_sk[i * (ldr + 1)] += shift; + last_shift = shift; + } + potrf_info = lapack::potrf(Uplo::Upper, n, R_sk, ldr); + if (potrf_info == 0) break; } - - // Cholesky factorization - if (lapack::potrf(Uplo::Upper, n, R_sk, ldr)) { - delete[] A_hat; - delete[] tau; - return 1; + this->n_chol_retries = (potrf_info == 0) ? attempt : (attempt > 0 ? attempt - 1 : 0); + this->chol_applied_shifts[0] = last_shift; + delete[] gram_backup; + if (potrf_info) { + delete[] A_hat; delete[] tau; return 1; } + if(this -> timing) potrf_t_stop = steady_clock::now(); - if(this -> timing) - potrf_t_stop = steady_clock::now(); - - // Obtain the output Q-factor (only if requested, timed separately and excluded from total) if (this->compute_Q) { - if(this -> timing) - q_t_start = steady_clock::now(); - - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, Diag::NonUnit, m, n, 1.0, R_sk, ldr, A, lda); - - if(this -> timing) - q_t_stop = steady_clock::now(); + if(this -> timing) q_t_start = steady_clock::now(); + blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, Diag::NonUnit, + m, n, 1.0, R_sk, ldr, A, lda); + if(this -> timing) q_t_stop = steady_clock::now(); } - if(this -> timing) - finalize_t_start = steady_clock::now(); - + if(this -> timing) finalize_t_start = steady_clock::now(); if (!this->orthogonalization) { - // Get the final R-factor - undoing the preconditioning - // R := R_chol * R_sk (returned by QR on sketch) - blas::trmm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, Diag::NonUnit, n, n, 1.0, A_hat, d, R_sk, ldr); + blas::trmm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, Diag::NonUnit, + n, n, 1.0, A_hat, d, R_sk, ldr); } + if(this -> timing) finalize_t_stop = steady_clock::now(); if(this -> timing) { - finalize_t_stop = steady_clock::now(); - total_t_stop = steady_clock::now(); - saso_t_dur = duration_cast(saso_t_stop - saso_t_start).count(); qr_t_dur = duration_cast(qr_t_stop - qr_t_start).count(); precond_t_dur = duration_cast(precond_t_stop - precond_t_start).count(); gram_t_dur = duration_cast(gram_t_stop - gram_t_start).count(); potrf_t_dur = duration_cast(potrf_t_stop - potrf_t_start).count(); finalize_t_dur = duration_cast(finalize_t_stop - finalize_t_start).count(); - - total_t_dur = duration_cast(total_t_stop - total_t_start).count(); - - // Subtract Q-factor computation time (excluded from total, matching CQRRT_linops test_mode) + total_t_dur = duration_cast(total_t_stop - total_t_start).count(); if (this->compute_Q) { q_t_dur = duration_cast(q_t_stop - q_t_start).count(); total_t_dur -= q_t_dur; } - long t_rest = total_t_dur - (saso_t_dur + qr_t_dur + precond_t_dur + gram_t_dur + potrf_t_dur + finalize_t_dur); - - // Fill the data vector (10 entries, matching CQRRT_linops indices) - // Index: 0=saso, 1=qr, 2=trtri(=0), 3=precond, 4=gram, 5=trmm_gram(=0), 6=potrf, 7=finalize, 8=rest, 9=total this -> times = {saso_t_dur, qr_t_dur, 0L, precond_t_dur, gram_t_dur, 0L, potrf_t_dur, finalize_t_dur, t_rest, total_t_dur}; @@ -282,7 +314,226 @@ int CQRRT::call( delete[] A_hat; delete[] tau; - return 0; } + + +// ============================================================================ +// CQRRT_linops: sketch-preconditioned Q-less Cholesky QR for abstract operators +// ============================================================================ +// +// Algorithm 4 from the collaborator's spec. Cannot modify the operator in place, +// so it forms R_sk explicitly and delegates to cholqr_primitive (which handles +// the precondition-inversion strategy via PCholQRPrecondMethod). +// +template +class CQRRT_linops { + public: + + bool timing; + bool test_mode; + + // Q-factor for test mode (only allocated if test_mode = true) + T* Q; + int64_t Q_rows; + int64_t Q_cols; + + // 11 entries (preserved for matlab plotter compatibility): + // [0] alloc, [1] sketch, [2] qr, [3] tri_inv, [4] fwd, [5] adj, [6] trsm_gram, + // [7] chol, [8] finalize, [9] rest, [10] total + std::vector times; + /// Total measured wall-clock (microseconds) of the last call(), or -1 if timing + /// was off. Every driver in this family packs the total as the LAST times[] entry, + /// but the entry COUNT differs per driver (6 / 11 / 15 / 18). Callers used to hard- + /// code that index (times[5], times[10], times[14], times[17]), so adding or + /// removing one slot silently wrote the wrong number into every CSV with no compile + /// error. Read the total through here instead. + long total_us() const { return times.empty() ? -1L : times.back(); } + + int64_t nnz; + CQRRTLinopPrecond precond_method; + T bqrrp_block_ratio; + int64_t block_size; + + // Adaptive-shift safety net (same as CholQR/CholQR2): the preconditioned + // Gram Cholesky's first attempt is always unshifted; only if potrf breaks + // down does the primitive seed the shift at eps*trace(G) and grow it x + // shift_growth. max_retries < 0 = unbounded, retry until PD. This lets + // CQRRT survive an ill-conditioned (e.g. single-precision) Gram instead + // of failing outright, matching the CholQR family. + int max_retries; + T shift_growth; + int n_chol_retries = 0; ///< shift retries used on the last call (0 = clean) + /// Shift record from the last call's preconditioned-Gram Cholesky: absolute + /// shift the successful potrf carried (0 = unshifted) and the Gram's trace. + T chol_applied_shifts[1] = {T(0)}; + T chol_gram_traces[1] = {T(0)}; + + CQRRT_linops( + bool time_subroutines, + T ep, + bool enable_test_mode = false + ) { + timing = time_subroutines; + (void)ep; // kept in the signature for call-site compatibility; unused + nnz = 4; // SASO nonzeros per column; paper uses 4 or 8 (2 causes sporadic spikes) + block_size = kDefaultGramBlockSize; + precond_method = PCholQRPrecondMethod::TRSM_IDENTITY; + bqrrp_block_ratio = (T)1.0; + max_retries = -1; // unbounded retries (no ceiling), as CholQR/CholQR2 + shift_growth = T(10); + test_mode = enable_test_mode; + Q = nullptr; + Q_rows = 0; + Q_cols = 0; + } + + ~CQRRT_linops() { + if (Q != nullptr) { + delete[] Q; + } + } + + template + int call( + GLO& A, + T* R, + int64_t ldr, + T d_factor, + RandBLAS::RNGState &state + ) { + steady_clock::time_point t0, t1, total_t_start, total_t_stop; + long alloc_dur = 0, saso_dur = 0, qr_dur = 0; + long precond_inv_dur = 0, fwd_dur = 0, adj_dur = 0, gemm_dur = 0; + long chol_dur = 0, finalize_dur = 0, q_dur = 0, total_dur = 0; + + if (this->timing) total_t_start = steady_clock::now(); + + int64_t m = A.n_rows; + int64_t n = A.n_cols; + // Input validation: d_factor < 1 gives d < n, which reads out of bounds + // in the lacpy of the upper n x n block below. + randlapack_require(m >= n) << "CQRRT_linops: operator must be tall (m=" << m << " < n=" << n << ")"; + randlapack_require(n >= 1) << "CQRRT_linops: n must be >= 1"; + randlapack_require(d_factor >= (T)1.0) << "CQRRT_linops: d_factor=" << d_factor << " must be >= 1.0"; + randlapack_require(ldr >= n) << "CQRRT_linops: ldr=" << ldr << " < n=" << n; + randlapack_require(R != nullptr) << "CQRRT_linops: R buffer is null"; + int64_t d = (int64_t)(d_factor * (T)n); + if (d < n) d = n; + int64_t b_eff = (this->block_size > 0 && this->block_size < n) + ? this->block_size : n; + + // ---- Allocations, PHASED ---- + // The sketch phase and the Gram/Cholesky phase have disjoint working sets, + // so allocating both up front sums their peaks. Phasing drops the peak + // from (d*n + n + n^2) + (3n^2 + (m+n)*b_eff + n) to + // max(d*n + n + n^2, 3n^2 + (m+n)*b_eff + n): sketch moment is + // A_hat(d*n) + tau(n) + P(n*n); Gram moment is P + R_pre + G (3*n*n) + // + A_temp(m*b_eff) + Z_buf(n*b_eff) + cholqr_primitive's O(n) + // diag_backup. See cqrrt_linops_analytical_kb in rl_memory_tracker.hh, + // the source of truth this comment must agree with. + if (this->timing) t0 = steady_clock::now(); + T* A_hat = new T[d * n]; + T* tau = new T[n]; + T* P = new T[n * n](); + T* R_pre = nullptr; // Gram-phase buffers: allocated after the sketch QR, + T* G = nullptr; // once A_hat and tau have been released. + T* A_temp = nullptr; + T* Z_buf = nullptr; + if (this->timing) { t1 = steady_clock::now(); alloc_dur = duration_cast(t1 - t0).count(); } + + // ---- Step 1: Sketch M^sk = S * A ---- + // (Sparse SASO only.) + if (this->timing) t0 = steady_clock::now(); + { + RandBLAS::SparseDist DS(d, m, this->nnz); + RandBLAS::SparseSkOp S(DS, state); + state = S.next_state; + RandBLAS::fill_sparse(S); + A(Side::Right, Layout::ColMajor, Op::NoTrans, Op::NoTrans, + d, n, m, (T)1.0, S, (T)0.0, A_hat, d); + } + if (this->timing) { t1 = steady_clock::now(); saso_dur = duration_cast(t1 - t0).count(); } + + // ---- Step 2: [~, R^sk] = qr(M^sk) ---- + if (this->timing) t0 = steady_clock::now(); + lapack::geqrf(d, n, A_hat, d, tau); + lapack::lacpy(MatrixType::Upper, n, n, A_hat, d, P, n); // R^sk = upper(A_hat); P's lower stays 0 + // The sketch and its Householder scalars are dead here: CQRRT never forms + // or applies the sketch Q, only R^sk, which now lives in P. Release them + // BEFORE the Gram-phase buffers exist. + delete[] A_hat; A_hat = nullptr; + delete[] tau; tau = nullptr; + if (this->timing) { t1 = steady_clock::now(); qr_dur = duration_cast(t1 - t0).count(); } + + // Gram/Cholesky working set, allocated now that the sketch is gone. + if (this->timing) t0 = steady_clock::now(); + R_pre = new T[n * n](); + G = new T[n * n](); + A_temp = new T[m * b_eff]; + Z_buf = new T[n * b_eff]; + if (this->timing) { t1 = steady_clock::now(); alloc_dur += duration_cast(t1 - t0).count(); } + + // ---- Step 3: PCholQR(A, P = R^sk) ---- + int info = cholqr_primitive( + A, P, R, ldr, + this->precond_method, + this->block_size, + this->bqrrp_block_ratio, + R_pre, G, A_temp, Z_buf, + &state, + precond_inv_dur, fwd_dur, adj_dur, gemm_dur, chol_dur, finalize_dur, + this->timing, + /*shift_factor=*/T(0), this->max_retries, this->shift_growth, &this->n_chol_retries, + &this->chol_applied_shifts[0], &this->chol_gram_traces[0]); + if (info != 0) { + // cholqr_primitive's status (retry exhaustion, singular + // preconditioner, non-finite shift) is flattened to 1 here; + // see stderr for the specific cause. + delete[] A_hat; delete[] tau; delete[] P; delete[] R_pre; + delete[] G; delete[] A_temp; delete[] Z_buf; + return 1; + } + + // ---- Test mode: materialize Q = A * R^{-1} ---- + if (this->test_mode) { + if (this->timing) t0 = steady_clock::now(); + + T* Q_buf = new T[m * n]; + RandLAPACK::materialize_Q_from_R(A, R, ldr, m, n, b_eff, Q_buf, m); + this->Q_rows = m; + this->Q_cols = n; + // The class owns Q (the destructor frees it), so release any buffer from a + // previous call() before taking ownership of this one. + delete[] this->Q; + this->Q = Q_buf; + + if (this->timing) { t1 = steady_clock::now(); q_dur = duration_cast(t1 - t0).count(); } + } + + // ---- Finalize timing ---- + if (this->timing) { + total_t_stop = steady_clock::now(); + total_dur = duration_cast(total_t_stop - total_t_start).count(); + total_dur -= q_dur; + + long rest_dur = total_dur - (alloc_dur + saso_dur + qr_dur + precond_inv_dur + + fwd_dur + adj_dur + gemm_dur + chol_dur + finalize_dur); + + this->times = {alloc_dur, saso_dur, qr_dur, precond_inv_dur, + fwd_dur, adj_dur, gemm_dur, + chol_dur, finalize_dur, rest_dur, total_dur}; + } + + delete[] A_hat; + delete[] tau; + delete[] P; + delete[] R_pre; + delete[] G; + delete[] A_temp; + delete[] Z_buf; + return 0; + } +}; + } // end namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_cqrrt_linops.hh b/RandLAPACK/drivers/rl_cqrrt_linops.hh deleted file mode 100644 index b61be25c1..000000000 --- a/RandLAPACK/drivers/rl_cqrrt_linops.hh +++ /dev/null @@ -1,449 +0,0 @@ -#pragma once - -#include "rl_util.hh" -#include "rl_blaspp.hh" -#include "rl_lapackpp.hh" -#include "rl_linops.hh" - -#include -#include -#include -#include - -using namespace std::chrono; - -namespace RandLAPACK { - -/// Sketch-preconditioned Cholesky QR for abstract linear operators. -/// -/// Linop analogue of CQRRT (rl_cqrrt.hh). Computes A = QR where A is any -/// type satisfying the LinearOperator concept. The algorithm sketches A to -/// obtain a preconditioner R_sk, then computes the Gram matrix of the -/// preconditioned operator A * R_sk^{-1} via linop calls, and factors it -/// with Cholesky. -/// -/// Unlike rl_cqrrt.hh (which overwrites A in place with TRSM), this class -/// cannot modify the operator directly. Instead, it computes R_sk^{-1} -/// explicitly and multiplies through the operator interface. -/// -/// The Q factor is not computed by default (Q-less factorization). When -/// test_mode is enabled, Q is materialized for verification. -/// -template -class CQRRT_linops { - public: - - bool timing; - bool test_mode; - T eps; - - // Q-factor for test mode (only allocated if test_mode = true) - T* Q; - int64_t Q_rows; - int64_t Q_cols; - - // 11 entries: alloc, sketch, qr, tri_inv, fwd, adj, trmm, chol, finalize, rest, total - // fwd = LinOp NoTrans: A * R_sk_inv (accumulated over blocks) - // adj = LinOp Trans: A^T * buf (accumulated over blocks) - // trmm = R_sk_inv^T * G (dense trmm, completes Gram) - std::vector times; - - // tuning SASOS - int64_t nnz; - - // If true, use a dense Gaussian sketching operator instead of sparse SASO. - // Dense sketches avoid potential issues with rank-deficient sketches but - // require O(d*m) storage and O(d*m*n) work for application. - bool use_dense_sketch; - - // Column-block size for the precondition + Gram computation. - // - // When block_size > 0, the two expensive linear operator calls: - // (1) A_pre = A * R_sk_inv (m × n) - // (2) R = A^T * A_pre (n × n) - // are fused into a column-block loop that processes b columns at a time: - // for each column block j of width b: - // buf (m × b) = A * R_sk_inv[:, j*b : (j+1)*b] - // R[:, j*b : (j+1)*b] = A^T * buf - // - // This reduces peak memory from O(m*n) to O(m*b), which is significant - // when m is large and n is moderate. The result is mathematically - // identical — each column block of R is: - // R[:, j_block] = A^T * (A * R_sk_inv[:, j_block]) - // which equals the corresponding columns of A^T * A * R_sk_inv. - // - // When block_size <= 0 or block_size >= n, the full m × n buffer is - // allocated and the original (non-blocked) path is used. - // - // When test_mode is enabled and blocking is active, the Q-factor - // computation (which needs the full m × n A_pre) is handled by - // recomputing A_pre = A * R_sk_inv after the Gram loop. This - // recomputation is outside the timing region, so it does not - // affect benchmark results. - int64_t block_size; - - CQRRT_linops( - bool time_subroutines, - T ep, - bool enable_test_mode = false - ) { - timing = time_subroutines; - eps = ep; - nnz = 2; - use_dense_sketch = false; - block_size = 0; - test_mode = enable_test_mode; - Q = nullptr; - Q_rows = 0; - Q_cols = 0; - } - - ~CQRRT_linops() { - if (Q != nullptr) { - delete[] Q; - } - } - - /// Computes the R-factor of the unpivoted QR factorization A = QR, - /// where Q is m-by-n and R is n-by-n upper triangular. - /// - /// Operates similarly to rl_cqrrt.hh, but accepts any type satisfying - /// the LinearOperator concept and returns a Q-less factorization - /// (Q is only computed when test_mode is enabled). - /// - /// Algorithm: - /// 1. Sketch: S*A (d x n), where d = d_factor * n - /// 2. QR of sketch: S*A = Q_sk * R_sk - /// 3. Precondition: A_pre = A * R_sk^{-1} - /// 4. Gram matrix: G = (R_sk^{-1})^T * A^T * A * R_sk^{-1} - /// 5. Cholesky: G = R_chol^T * R_chol - /// 6. Final R = R_chol * R_sk - /// - /// @note This algorithm expects A to be full-rank (rank = n). Rank-deficient inputs may result - /// in loss of orthogonality in the Q-factor (when test_mode=true) and numerical instability - /// in the R-factor. - /// - /// @param[in] A - /// The m-by-n linear operator (m and n read from A.n_rows, A.n_cols). - /// - /// @param[out] R - /// Pre-allocated n-by-n buffer. On exit, stores the upper-triangular - /// R factor. Zero entries are not compressed. - /// - /// @param[in] ldr - /// Leading dimension of R. - /// - /// @param[in] d_factor - /// Sketch embedding factor. The sketch dimension is d = d_factor * n. - /// Typically d_factor >= 1; larger values improve numerical stability. - /// - /// @param[in,out] state - /// RNG state for sketching operator generation. Advanced on exit. - /// - /// @return = 0: successful exit - template - int call( - GLO& A, - T* R, - int64_t ldr, - T d_factor, - RandBLAS::RNGState &state - ) { - ///--------------------TIMING VARS--------------------/ - steady_clock::time_point t_start, t_stop; - steady_clock::time_point alloc_t_start, alloc_t_stop; - steady_clock::time_point saso_t_start, saso_t_stop; - steady_clock::time_point qr_t_start, qr_t_stop; - steady_clock::time_point trtri_t_start, trtri_t_stop; - steady_clock::time_point trmm_gram_t_start, trmm_gram_t_stop; - steady_clock::time_point potrf_t_start, potrf_t_stop; - steady_clock::time_point finalize_t_start, finalize_t_stop; - steady_clock::time_point total_t_start, total_t_stop; - steady_clock::time_point q_t_start, q_t_stop; - long alloc_t_dur = 0; - long saso_t_dur = 0; - long qr_t_dur = 0; - long trtri_t_dur = 0; - long fwd_t_dur = 0; - long adj_t_dur = 0; - long trmm_gram_t_dur = 0; - long potrf_t_dur = 0; - long finalize_t_dur = 0; - long total_t_dur = 0; - long q_t_dur = 0; - - if(this -> timing) - total_t_start = steady_clock::now(); - - int64_t m = A.n_rows; - int64_t n = A.n_cols; - - int64_t d = d_factor * n; - - if(this -> timing) - alloc_t_start = steady_clock::now(); - - // No zero-init: first use is with beta=0.0 which overwrites all elements - T* A_hat = new T[d * n]; - // No zero-init: geqrf writes the output - T* tau = new T[n]; - - if(this -> timing) { - alloc_t_stop = steady_clock::now(); - saso_t_start = steady_clock::now(); - } - - /// Generate and apply the sketching operator to the linear operator. - // Side::Right means the operator A is on the right side: C = op(S) * op(A) - if (this->use_dense_sketch) { - // Dense Gaussian sketch: allocate d x m buffer, fill with Gaussian entries. - RandBLAS::DenseDist DD(d, m); - RandBLAS::DenseSkOp S(DD, state); - state = S.next_state; - RandBLAS::fill_dense(S); - A(Side::Right, Layout::ColMajor, Op::NoTrans, Op::NoTrans, d, n, m, (T)1.0, S, (T)0.0, A_hat, d); - } else { - // Sparse SASO sketch: uses nnz nonzeros per column. - RandBLAS::SparseDist DS(d, m, this->nnz); - RandBLAS::SparseSkOp S(DS, state); - state = S.next_state; - RandBLAS::fill_sparse(S); - A(Side::Right, Layout::ColMajor, Op::NoTrans, Op::NoTrans, d, n, m, (T)1.0, S, (T)0.0, A_hat, d); - } - - if(this -> timing) { - saso_t_stop = steady_clock::now(); - qr_t_start = steady_clock::now(); - } - - /// Performing QR on a sketch - lapack::geqrf(d, n, A_hat, d, tau); - - if(this -> timing) - qr_t_stop = steady_clock::now(); - - // Compute R_sk^{-1} via TRSM with identity (since A may not be dense, - // we cannot apply TRSM directly to the operator as in rl_cqrrt.hh). - if(this -> timing) - trtri_t_start = steady_clock::now(); - - // Instead of doing TRTRI to find R_sk_inv, we do TRSM with an identity, since trtri is not optimized in MKL - T* Eye = new T[n * n](); - RandLAPACK::util::eye(n, n, Eye); - if (!RandLAPACK::util::diag_is_nonzero(n, A_hat, d)) { - delete[] A_hat; - delete[] tau; - delete[] Eye; - return 1; - } - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, Diag::NonUnit, n, n, 1.0, A_hat, d, Eye, n); - if (n > 1) { - // Clear the below-diagonal - lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &Eye[1], n); - } - T* R_sk_inv = Eye; - - if(this -> timing) { - trtri_t_stop = steady_clock::now(); - } - - // Gram computation: R = (R_sk_inv)^T * A^T * A * R_sk_inv - // The (R_sk_inv)^T left-multiply is handled later via TRMM. - // Here we compute: R = A^T * (A * R_sk_inv). - - // Determine effective block width. - // block_size <= 0 or >= n means "no blocking" (full width). - int64_t b_eff = (this->block_size > 0 && this->block_size < n) - ? this->block_size : n; - - // A_pre buffer: - // Full path: m × n (kept alive for Q-factor in test_mode) - // Block path: m × b_eff (temporary, freed after Gram loop; - // if test_mode, a full m × n buffer is - // allocated later for Q computation) - // No zero-init: first use is with beta=0.0 which overwrites all elements - T* A_pre = new T[m * b_eff]; - - if (b_eff == n) { - // --- Full materialization path (original) --- - - // Step 1: A_pre (m × n) = A * R_sk_inv (m × n × n) - // A is m × n, R_sk_inv is n × n, A_pre is m × n. - // Side::Left: C = alpha * op(A) * op(B) + beta * C - // where op(A) = A (m × n), op(B) = R_sk_inv (n × n), C = A_pre (m × n). - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, m, n, n, (T)1.0, R_sk_inv, n, (T)0.0, A_pre, m); - if(this->timing) { t_stop = steady_clock::now(); fwd_t_dur = duration_cast(t_stop - t_start).count(); } - - // Step 2: R (n × n) = A^T * A_pre (n × m × m × n = n × n) - // Since SYRK is not defined for non-dense operators, we use - // an explicit A^T * A_pre to form the Gram matrix. - // op(A) = A^T (n × m), op(B) = A_pre (m × n), C = R (n × n). - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, n, n, m, (T)1.0, A_pre, m, (T)0.0, R, ldr); - if(this->timing) { t_stop = steady_clock::now(); adj_t_dur = duration_cast(t_stop - t_start).count(); } - } else { - // --- Column-block processing path (memory-efficient) --- - // - // Process n columns in blocks of width b_eff. - // The last block may be narrower if n is not divisible by b_eff. - // - // For each block starting at column j with width b_j: - // (1) buf (m × b_j) = A * R_sk_inv[:, j : j+b_j] - // R_sk_inv is n × n in ColMajor with ld = n. - // Column j starts at R_sk_inv + j * n. - // We multiply A (m × n) by this n × b_j slice. - // - // (2) R[:, j : j+b_j] (n × b_j) = A^T * buf - // A^T is n × m, buf is m × b_j, result is n × b_j. - // R is n × n with ld = ldr. - // Column j starts at R + j * ldr. - // - // Total FLOPs are the same as the full path; only memory differs. - - long fwd_accum = 0, adj_accum = 0; - for (int64_t j = 0; j < n; j += b_eff) { - int64_t b_j = std::min(b_eff, n - j); - - // (1) buf = A * R_sk_inv[:, j : j+b_j] (NoTrans = fwd) - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, b_j, n, (T)1.0, R_sk_inv + j * n, n, (T)0.0, A_pre, m); - if(this->timing) { t_stop = steady_clock::now(); fwd_accum += duration_cast(t_stop - t_start).count(); } - - // (2) R[:, j : j+b_j] = A^T * buf (Trans = adj) - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, - n, b_j, m, (T)1.0, A_pre, m, (T)0.0, R + j * ldr, ldr); - if(this->timing) { t_stop = steady_clock::now(); adj_accum += duration_cast(t_stop - t_start).count(); } - } - - if(this -> timing) { - fwd_t_dur = fwd_accum; - adj_t_dur = adj_accum; - } - } - - if(this -> timing) { - trmm_gram_t_start = steady_clock::now(); - } - - // (R_sk_inv)^T * (A^T * A_pre) - blas::trmm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::Trans, Diag::NonUnit, n, n, (T) 1.0, R_sk_inv, n, R, ldr); - - if(this -> timing) { - trmm_gram_t_stop = steady_clock::now(); - potrf_t_start = steady_clock::now(); - } - - // Cholesky factorization (only reads/writes upper triangle) - if (lapack::potrf(Uplo::Upper, n, R, ldr)) { - delete[] A_hat; - delete[] tau; - delete[] R_sk_inv; - delete[] A_pre; - return 1; - } - - if(this -> timing) - potrf_t_stop = steady_clock::now(); - - // Compute Q-factor if test mode is enabled (NOT included in cholqr timing) - if(this->test_mode) { - if(this->timing) - q_t_start = steady_clock::now(); - - if (b_eff < n) { - // Column-block Gram was used: A_pre is only m × b_eff, - // too small for Q. Recompute A_pre = A * R_sk_inv in - // full. This is outside the timing region, so the extra - // operator application does not affect benchmark results. - delete[] A_pre; - // No zero-init: beta=0.0 overwrites all elements - A_pre = new T[m * n]; - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, n, n, (T)1.0, R_sk_inv, n, (T)0.0, A_pre, m); - } - - // Reuse A_pre storage for Q (Q = A_pre * R_chol^{-1}) - this->Q_rows = m; - this->Q_cols = n; - this->Q = A_pre; // Take ownership of A_pre buffer - - // Solve Q * R_chol = A_pre for Q (R_chol is upper triangular from Cholesky) - // Q = A * (R_chol * R_sk)^{-1} = A * R_sk^{-1} * R_chol^{-1} = A_pre * R_chol^{-1} - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, m, n, (T)1.0, R, ldr, this->Q, m); - - if(this->timing) - q_t_stop = steady_clock::now(); - } - - // Zero out strictly lower triangle of R before final trmm - // trmm expects R to be upper triangular on input (it preserves triangular structure) - // The lower triangle may contain garbage from the Gram matrix computation - // Use laset with beta=1.0 to preserve diagonal while zeroing strictly lower triangle - if (n > 1) { - lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &R[1], ldr); - } - - if(this -> timing) - finalize_t_start = steady_clock::now(); - - // Get the final R-factor - undoing the preconditioning - // R := R_chol * R_sk, where R_sk is the upper triangle of A_hat - // trmm with Uplo::Upper only reads upper triangle of A_hat (can use ld=d directly) - // and expects R to be upper triangular on input - blas::trmm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, Diag::NonUnit, n, n, 1.0, A_hat, d, R, ldr); - - if(this -> timing) - finalize_t_stop = steady_clock::now(); - - if(this -> timing) { - // Stop timing BEFORE cleanup operations to exclude deallocation costs. - // Note: First iteration may show inflated times due to cold cache effects. - total_t_stop = steady_clock::now(); - - alloc_t_dur = duration_cast(alloc_t_stop - alloc_t_start).count(); - saso_t_dur = duration_cast(saso_t_stop - saso_t_start).count(); - qr_t_dur = duration_cast(qr_t_stop - qr_t_start).count(); - trtri_t_dur = duration_cast(trtri_t_stop - trtri_t_start).count(); - // fwd_t_dur and adj_t_dur already set (in both full and blocked paths) - trmm_gram_t_dur = duration_cast(trmm_gram_t_stop - trmm_gram_t_start).count(); - potrf_t_dur = duration_cast(potrf_t_stop - potrf_t_start).count(); - finalize_t_dur = duration_cast(finalize_t_stop - finalize_t_start).count(); - - total_t_dur = duration_cast(total_t_stop - total_t_start).count(); - - // Subtract Q-factor computation time if in test mode - if(this->test_mode) { - q_t_dur = duration_cast(q_t_stop - q_t_start).count(); - total_t_dur -= q_t_dur; - } - - long t_rest = total_t_dur - (alloc_t_dur + saso_t_dur + qr_t_dur + trtri_t_dur + fwd_t_dur + - adj_t_dur + trmm_gram_t_dur + potrf_t_dur + finalize_t_dur); - - // Fill the data vector (11 entries) - // Index: 0=alloc, 1=sketch, 2=qr, 3=tri_inv, 4=fwd, 5=adj, 6=trmm, 7=chol, 8=finalize, 9=rest, 10=total - this -> times = {alloc_t_dur, saso_t_dur, qr_t_dur, trtri_t_dur, fwd_t_dur, - adj_t_dur, trmm_gram_t_dur, potrf_t_dur, finalize_t_dur, - t_rest, total_t_dur}; - } - - // Cleanup - now outside the timing region to avoid timing artifacts - delete[] A_hat; - delete[] tau; - delete[] R_sk_inv; - - // Only delete A_pre if not in test mode (otherwise Q owns it). - // When test_mode + blocking: the small block buffer was freed - // and replaced with a full m × n buffer in the Q section above. - if(!this->test_mode) { - delete[] A_pre; - } - - return 0; - } -}; -} // end namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_iter_refine_lsq.hh b/RandLAPACK/drivers/rl_iter_refine_lsq.hh new file mode 100644 index 000000000..8f26a535e --- /dev/null +++ b/RandLAPACK/drivers/rl_iter_refine_lsq.hh @@ -0,0 +1,280 @@ +#pragma once + +// Public API: IterRefineLSQ: Q-less, sketch-and-precondition iterative-refinement +// least-squares solver. +// +// Solves min_x ||b - J x||_2 for a tall LinearOperator J using a precomputed +// triangular preconditioner R (e.g., the R-factor from CQRRT_linops on J or on +// a sketch SJ). R is treated as a right preconditioner on the normal equations. +// +// This class is a thin adapter over the shared restarted engine +// restarted_pcg_ne (rl_restarted_pcg_ne.hh); it adds the historical field +// names, the per-step diagnostic vectors, and the cold-start policy. +// +// RESTART PACING. Each round's inner CG stops after its residual has dropped +// by the factor `round_drop` (default 1e-4) relative to the round's own +// right-hand side, rather than grinding to a fixed tiny tolerance: only the +// between-round recomputation of the true residual injects new information, +// so deep inner solves polish a stale right-hand side. `inner_tol` survives +// as the ABSOLUTE inner target: a round whose normal-equation residual has +// already fallen to inner_tol * ||g_0|| terminates immediately (the "CG +// still stops once below the target" guard). Set round_drop <= 0 to restore +// the legacy fixed-tolerance rounds. +// +// Reference: E. N. Epperly, M. Meier, and Y. Nakatsukasa, +// "Fast randomized least-squares solvers can be just as accurate and stable +// as classical direct solvers," arXiv:2406.03468v3 (2025), Algorithm 1 +// + Theorem 6.1 (master theorem on backward stability of two-step IR). + +#include "rl_blaspp.hh" +#include "rl_exceptions.hh" +#include "rl_pcg_inner.hh" +#include "rl_restarted_pcg_ne.hh" +#include "../linops/rl_concepts.hh" + +#include +#include +#include +#include +#include +#include + + +namespace RandLAPACK { + + +/*********************************************************/ +/* */ +/* IterRefineLSQ */ +/* */ +/*********************************************************/ + +/// @brief Iterative-refinement least-squares solver with right preconditioner R. +/// +/// Solves min_x ||b - J x||_2 by up to `n_refine_steps` rounds of +/// +/// r_i ← b - J x_i (true residual, working precision) +/// c_i ← R^{-T} (J^T r_i) +/// z_i ← CG on M z = c_i with M = R^{-T} J^T J R^{-1}, +/// stopped after a `round_drop` residual +/// drop (or at the inner_tol floor) +/// x_{i+1} ← x_i + R^{-1} z_i +/// +/// exiting early once ||b - J x|| / ||b|| <= outer_tol. Executed by the shared +/// engine restarted_pcg_ne; see the file header for the pacing rationale. +// InnerCGStatus lives in rl_pcg_inner.hh (same name, same namespace, same +// values; callers are unaffected). + +template +struct IterRefineLSQ { + // ------------- Configuration ------------- + /// ABSOLUTE inner-CG target, relative to the initial normal-equation + /// right-hand side ||R^{-T} J^T b||: a round whose NE residual is already + /// below inner_tol * ||g_0|| stops immediately. In legacy mode + /// (round_drop <= 0) this is instead the per-round relative tolerance, + /// the legacy contract. + T inner_tol; + /// Hard cap on inner CG iterations per round. The TOTAL budget across all + /// rounds is max_inner_iters * n_refine_steps. + int max_inner_iters; + /// STAGNATION EXIT. Stop the inner CG when its residual has not improved + /// significantly for `inner_stag_window` consecutive iterations, and + /// return the BEST iterate seen rather than the last one (on an + /// ill-conditioned preconditioner, raising the iteration cap does not + /// improve the best residual but does let the last iterate drift worse). + /// Set `inner_stag_window <= 0` to disable. A converging solve is + /// unaffected: it returns Converged before the window elapses. + int inner_stag_window; + /// Relative residual drop that counts as progress for the stagnation test + /// (default 1e-3, i.e. the residual must fall by at least 0.1%). + T inner_stag_rel_improve; + /// RESTART PACING: per-round relative residual drop at which the inner CG + /// stops and control returns to the outer loop for a true-residual + /// restart. Default 1e-4. <= 0 restores the legacy contract (each round + /// runs to the fixed inner_tol). + T round_drop; + /// Maximum outer rounds (default benchmark setting: 20). With outer_tol + /// enabled the loop reads "refine until done, capped at n_refine_steps"; + /// well-preconditioned methods exit after a few rounds, and the cap gives + /// weakly-preconditioned configurations room to keep descending. + int n_refine_steps; + /// OUTER EARLY EXIT: stop refining once the TRUE residual ||b - Jx|| / ||b|| + /// is at or below this value, checked between rounds where the residual is + /// recomputed anyway. 0 (the default) disables THIS check only; the run can + /// still end before n_refine_steps rounds via the engine's LS-floor + /// stagnation exit (see outer_stag_window) or a terminal inner-CG condition. + T outer_tol; + /// Consecutive rounds without significant true-residual improvement that end + /// the run as an LS-floor exit (forwarded to restarted_pcg_ne; see its + /// documentation). <= 0 disables the floor exit. Decoupled from + /// inner_stag_window: the two mechanisms are independent. + int outer_stag_window = 2; + /// Optional initial iterate to refine (length n), or nullptr for the cold + /// start that every Q-less method uses. Lets another solver's answer + /// (Blendenpik's) be handed to refinement, separating preconditioner + /// quality from solver structure in the benchmark suite. + const T* warm_x0 = nullptr; + /// Enable per-step / per-substep timing breakdown. + bool timing; + /// Print convergence info to stdout. + bool verbose; + + // ------------- Outputs (filled by call) ------------- + /// Number of outer rounds actually executed. + int outer_iters_done; + /// CG iteration counts for each round. + std::vector inner_iters_per_step; + /// Exit condition of each round's inner CG (see InnerCGStatus). + std::vector inner_status_per_step; + /// Relative CG residual ||M z - c|| / ||c|| at exit, per round. + std::vector inner_relres_per_step; + /// Smallest relative CG residual seen during that round, and the iteration at + /// which it occurred. Together with inner_relres_per_step these separate the two + /// ways a solve can burn its budget: + /// best_iter << iters and best ~= final -> converged then STAGNATED (the + /// tolerance is below the attainable floor; more iterations cannot help) + /// best_iter ~= iters -> still descending at the cap (the + /// preconditioner is weak; more iterations would help) + std::vector inner_best_relres_per_step; + std::vector inner_best_iter_per_step; + /// True LS relative residual after each round (engine ls_relres, kept for + /// the per-round campaign sidecar records). + std::vector ls_relres_per_step; + /// Final relative residual ||b - J x|| / ||b|| (or ||b - J x|| if ||b|| == 0). + T final_residual_norm; + /// The engine's exit status, verbatim (see restarted_pcg_ne @returns: + /// 0 tol met, 1 budget, 2 breakdown, 3 round budget, 4 LS floor). Recorded + /// so callers can report WHY a run ended, not just whether it converged. + int engine_status = 1; + /// Total inner CG iterations summed over all rounds, for callers that report a + /// single iteration count. + int inner_iters_total() const { + int t = 0; for (int v : inner_iters_per_step) t += v; return t; + } + /// Per-substep wall-clock breakdown (microseconds), populated when timing == true. + /// Entries: [0]=outer_total, [1]=inner_cg_total, [2]=trsm_total, + /// [3]=fwd_total, [4]=adj_total, [5]=other. + /// Slot order (total first, trsm before fwd/adj) intentionally differs from + /// the engine's times[] (rl_restarted_pcg_ne.hh: [fwd, adj, trsm, total]); + /// this struct predates the engine unification. Benchmarks read both by + /// name, not by matching index, so the divergence is safe to keep. + std::vector times; + + IterRefineLSQ(T tol = std::pow(std::numeric_limits::epsilon(), (T)0.85), + int max_inner = 200, + int n_steps = 2, + bool timing_on = false, + bool verbose_on = false) + : inner_tol(tol), + max_inner_iters(max_inner), + inner_stag_window(20), + inner_stag_rel_improve((T)1e-3), + round_drop((T)1e-4), + n_refine_steps(n_steps), + outer_tol((T)0), + timing(timing_on), + verbose(verbose_on), + outer_iters_done(0), + final_residual_norm((T)0) + {} + + /// @brief Solve min ||b - J x||_2 with right preconditioner R. + /// + /// @tparam J_LO A LinearOperator type (must satisfy linops::LinearOperator). + /// + /// @param J Forward operator (m × n, m >= n; n_rows == m, n_cols == n). + /// @param R n × n upper triangular ColMajor; leading dim ldr >= n. + /// @param ldr Leading dimension of R. + /// @param b Right-hand side, length m. + /// @param m Number of rows of J / length of b. + /// @param x Solution buffer, length n. Incoming content is ignored: the + /// start is cold (the policy for Q-less methods) unless + /// warm_x0 is set, in which case THAT iterate is refined + /// (the Blendenpik handoff). + /// @param n Number of columns of J / length of x. + /// + /// @returns 0 on success; nonzero on inner-CG breakdown. + template + int call(J_LO& J, const T* R, int64_t ldr, + const T* b, int64_t m, T* x, int64_t n) + { + randlapack_require(n_refine_steps >= 1) + << "IterRefineLSQ: n_refine_steps must be >= 1"; + randlapack_require(max_inner_iters >= 1) + << "IterRefineLSQ: max_inner_iters must be >= 1"; + + using clock = std::chrono::steady_clock; + auto t_start = clock::now(); + + // Cold start unless the caller supplied warm_x0: no + // zero-fill needed here, restarted_pcg_ne unconditionally overwrites x + // (cold start or warm_x0 refinement, both handled inside the engine). + + // Legacy mode (round_drop <= 0): rounds run to the fixed inner_tol + // relative to their own right-hand side, no absolute guard. + const bool paced = (round_drop > (T)0); + T drop = paced ? round_drop : inner_tol; + T abs_guard = paced ? inner_tol : (T)0; + + PCGRoundHistory hist; + int iters_total = 0, rounds = 0; + T final_rel = (T)0; + long times4[4] = {0, 0, 0, 0}; + + int st = restarted_pcg_ne(J, m, n, R, ldr, b, x, + /*tol=*/outer_tol, + /*max_iters=*/max_inner_iters * n_refine_steps, + iters_total, + /*restart_maxit=*/max_inner_iters, + /*restart_drop=*/drop, + /*max_restarts=*/n_refine_steps - 1, + &rounds, + timing ? times4 : nullptr, + &final_rel, + inner_stag_window, inner_stag_rel_improve, + abs_guard, &hist, warm_x0, outer_stag_window); + engine_status = st; + + // Republish the engine's per-round records under the historical names. + inner_iters_per_step = hist.iters; + inner_status_per_step = hist.status; + inner_relres_per_step = hist.relres; + inner_best_relres_per_step = hist.best_relres; + inner_best_iter_per_step = hist.best_iter; + ls_relres_per_step = hist.ls_relres; + outer_iters_done = rounds; + final_residual_norm = final_rel; + + if (verbose) { + static const char* kNames[] = {"converged", "HIT CAP", "breakdown", "STAGNATED"}; + for (size_t s = 0; s < hist.iters.size(); ++s) { + std::printf("[IR-LSQ] round %zu: inner CG %s after %d iters, " + "relres=%.4e (best %.4e at iter %d); LS relres %.4e\n", + s, kNames[hist.status[s]], hist.iters[s], + (double)hist.relres[s], (double)hist.best_relres[s], + hist.best_iter[s], (double)hist.ls_relres[s]); + } + } + + if (timing) { + long total = std::chrono::duration_cast( + clock::now() - t_start).count(); + // Non-overlapping [outer_total, inner_cg_total, trsm, fwd, adj, other]: + // op totals are all-inclusive; "other" subtracts the kernel wallclock + // and the outer-only op time so nothing is counted twice. + long op_outer = (times4[0] - hist.t_fwd_inner_us) + + (times4[1] - hist.t_adj_inner_us) + + (times4[2] - hist.t_trsm_inner_us); + long other = total - hist.t_inner_us - op_outer; + if (other < 0) other = 0; + times = {total, hist.t_inner_us, times4[2], times4[0], times4[1], other}; + } + + // Historical contract: a capped or stagnated solve is not an error + // return; only a CG breakdown is. + return (st == 2) ? 1 : 0; + } +}; + + +} // namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_lsqr.hh b/RandLAPACK/drivers/rl_lsqr.hh new file mode 100644 index 000000000..03edcf0bf --- /dev/null +++ b/RandLAPACK/drivers/rl_lsqr.hh @@ -0,0 +1,215 @@ +#pragma once + +// Public API: lsqr: matrix-free LSQR (Paige & Saunders) least-squares solver, +// with an optional upper-triangular RIGHT preconditioner R. +// +// Solves min_x ||b - A x||_2 for a tall LinearOperator A (m x n, m >= n) using +// the Golub-Kahan bidiagonalization LSQR recurrence. Only matrix-vector products +// A*v and A^T*u are used (via the LinearOperator), so A is never materialized. +// +// If R (upper triangular, n x n, ColMajor) is supplied, LSQR is run on the +// right-preconditioned operator à = A R^{-1} (so it solves min_y ||b - à y||), +// and the returned x is R^{-1} y. This is exactly the Blendenpik use case: with a +// sketch-QR preconditioner R the operator à is nearly orthonormal (kappa(Ã) ~ 1), +// so LSQR converges in O(log(1/eps)) iterations independent of kappa(A). Pass +// R = nullptr for the unpreconditioned solve. +// +// Reference: C. C. Paige and M. A. Saunders, "LSQR: An algorithm for sparse +// linear equations and sparse least squares", ACM TOMS 8(1), 1982. Stopping +// tests S1 (||r||/||b|| <= btol) and S2 (||A^T r||/(||A|| ||r||) <= atol). + +#include "rl_blaspp.hh" +#include "rl_blas2_threads.hh" +#include "../linops/rl_concepts.hh" + +#include +#include +#include + + +namespace RandLAPACK { + + +/// @brief Matrix-free LSQR for min ||b - A x||, optional right preconditioner R. +/// +/// @param[in] A tall LinearOperator (m x n), applied as A*v and A^T*u. +/// @param[in] m,n dimensions (m >= n). +/// @param[in] R upper-triangular right preconditioner (n x n, ColMajor) +/// or nullptr for none. Must be nonsingular when supplied. +/// @param[in] ldr leading dimension of R. +/// @param[in] b right-hand side (length m). +/// @param[out] x solution (length n). +/// @param[in] atol tolerance for the ||A^T r|| stopping test (S2). +/// @param[in] btol tolerance for the ||r||/||b|| stopping test (S1). +/// @param[in] max_iters iteration cap. +/// @param[out] iters_done number of iterations actually run. +/// @param[out] times optional [fwd_us, adj_us, trsm_us, total_us] (may be nullptr). +/// @param[out] final_relres optional: the solver's own ||b - à y|| / ||b|| at +/// termination (the Paige-Saunders S1 estimate). For a right +/// preconditioner this equals ||b - A x|| / ||b|| since à y = A x. +/// @param[out] stop_test optional: WHICH test ended the run: 1 = S1 (residual met +/// btol), 2 = S2 (normal-equation test met atol; fires at the +/// LS floor even when the residual is far above btol), 0 = the +/// iteration cap. Callers comparing convergence flags against +/// restarted_pcg_ne (whose only success test is the true LS +/// residual) need this to tell S2 successes apart. +/// Polarity note: here 0 means the cap was hit (no stopping +/// test fired); in the engine's status codes (restarted_pcg_ne, +/// IterRefineLSQ::engine_status) 0 means success. Do not +/// compare the two as if they shared a convention. +/// @returns 0 if a stopping test was met; 1 if the iteration cap was hit. +template +int lsqr( + GLO& A, int64_t m, int64_t n, + const T* R, int64_t ldr, + const T* b, T* x, + T atol, T btol, int max_iters, + int& iters_done, + long* times = nullptr, + T* final_relres = nullptr, + int* stop_test = nullptr) +{ + using clock = std::chrono::steady_clock; + using std::chrono::duration_cast; + using std::chrono::microseconds; + // One width for the whole solve (see SolveWidthScope in rl_blas2_threads.hh). + SolveWidthScope solve_scope(n); + long t_fwd = 0, t_adj = 0, t_trsm = 0; + auto total_start = clock::now(); + if (stop_test) *stop_test = 0; + + const bool prec = (R != nullptr); + + // Workspaces (raw T*, freed before every return). + T* u = new T[m](); // left bidiag vector (length m) + T* v = new T[n](); // right bidiag vector (length n) + T* w = new T[n](); // update direction (length n) + T* av = new T[m](); // holds à v (length m) + T* atu = new T[n](); // holds Ã^T u (length n) + T* sc = new T[n](); // trsm scratch (length n) + auto cleanup = [&]() { delete[] u; delete[] v; delete[] w; delete[] av; delete[] atu; delete[] sc; }; + + // à v = A (R^{-1} v) (out has length m) + auto apply_Atilde = [&](const T* vin, T* out) { + const T* fwd_in = vin; + if (prec) { + std::copy(vin, vin + n, sc); // sc = v + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); // cap threads: see rl_blas2_threads.hh + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, n, R, ldr, sc, 1); + } + t_trsm += duration_cast(clock::now() - ts).count(); + fwd_in = sc; // sc = R^{-1} v + } + auto tf = clock::now(); + A(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, fwd_in, n, (T)0.0, out, m); + t_fwd += duration_cast(clock::now() - tf).count(); + }; + + // Ã^T u = R^{-T} (A^T u) (out has length n) + auto apply_AtildeT = [&](const T* uin, T* out) { + auto ta = clock::now(); + A(blas::Side::Left, blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + n, 1, m, (T)1.0, uin, m, (T)0.0, out, n); + t_adj += duration_cast(clock::now() - ta).count(); + if (prec) { + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); // cap threads: see rl_blas2_threads.hh + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::Trans, blas::Diag::NonUnit, n, R, ldr, out, 1); + } + t_trsm += duration_cast(clock::now() - ts).count(); + } + }; + + // ---- Bidiagonalization init: beta u = b, alpha v = Ã^T u ---- + std::copy(b, b + m, u); + T beta = blas::nrm2(m, u, 1); + T bnorm = beta; + for (int64_t i = 0; i < n; ++i) x[i] = (T)0; + if (beta == (T)0) { iters_done = 0; cleanup(); if (times) { times[0]=t_fwd; times[1]=t_adj; times[2]=t_trsm; times[3]=duration_cast(clock::now()-total_start).count(); } if (final_relres) *final_relres = (T)0; if (stop_test) *stop_test = 1; return 0; } + blas::scal(m, (T)1.0 / beta, u, 1); + + apply_AtildeT(u, v); + T alpha = blas::nrm2(n, v, 1); + if (alpha == (T)0) { + // b (mod preconditioning) is orthogonal to range(A): x = 0 is already + // the LS minimizer. Without this the next step divides by alpha, + // filling x with NaN that defeats both stop tests until the cap. + iters_done = 0; + cleanup(); + if (times) { times[0]=t_fwd; times[1]=t_adj; times[2]=t_trsm; times[3]=duration_cast(clock::now()-total_start).count(); } + if (final_relres) *final_relres = (T)1; + if (stop_test) *stop_test = 2; + return 0; + } + blas::scal(n, (T)1.0 / alpha, v, 1); + std::copy(v, v + n, w); + + T phibar = beta, rhobar = alpha; + // Frobenius-norm accumulator for the S2 test; starts from the init-phase + // alpha_1 (previously omitted, slightly understating the estimate). + T anorm2 = alpha * alpha; + iters_done = 0; + int status = 1; + + for (int it = 1; it <= max_iters; ++it) { + // u ← à v - alpha u ; beta = ||u|| ; normalize + apply_Atilde(v, av); + blas::scal(m, -alpha, u, 1); + blas::axpy(m, (T)1.0, av, 1, u, 1); + beta = blas::nrm2(m, u, 1); + if (beta > (T)0) blas::scal(m, (T)1.0 / beta, u, 1); + + // v ← Ã^T u - beta v ; alpha = ||v|| ; normalize + apply_AtildeT(u, atu); + blas::scal(n, -beta, v, 1); + blas::axpy(n, (T)1.0, atu, 1, v, 1); + alpha = blas::nrm2(n, v, 1); + if (alpha > (T)0) blas::scal(n, (T)1.0 / alpha, v, 1); + + // Orthogonal transformation (plane rotation) + T rho = std::hypot(rhobar, beta); + T c = rhobar / rho; + T s = beta / rho; + T theta = s * alpha; + rhobar = -c * alpha; + T phi = c * phibar; + phibar = s * phibar; + + // y ← y + (phi/rho) w ; w ← v - (theta/rho) w + blas::axpy(n, phi / rho, w, 1, x, 1); + blas::scal(n, -theta / rho, w, 1); + blas::axpy(n, (T)1.0, v, 1, w, 1); + + // Stopping tests (Paige-Saunders estimates) + anorm2 += alpha * alpha + beta * beta; + T anorm = std::sqrt(anorm2); + T rnorm = phibar; // ||b - à y|| + T arnorm = phibar * alpha * std::abs(c); // ||Ã^T r|| + iters_done = it; + if (rnorm <= btol * bnorm) { status = 0; if (stop_test) *stop_test = 1; break; } + if (anorm * rnorm > (T)0 && arnorm <= atol * anorm * rnorm) { status = 0; if (stop_test) *stop_test = 2; break; } + } + + // Undo the preconditioner: x = R^{-1} y (y currently in x). + if (prec) { + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); // cap threads: see rl_blas2_threads.hh + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, n, R, ldr, x, 1); + } + t_trsm += duration_cast(clock::now() - ts).count(); + } + + if (times) { times[0]=t_fwd; times[1]=t_adj; times[2]=t_trsm; times[3]=duration_cast(clock::now()-total_start).count(); } + // phibar holds the last ||b - à y|| estimate; bnorm is ||b||. + if (final_relres) *final_relres = (bnorm > (T)0) ? (phibar / bnorm) : (T)0; + cleanup(); + return status; +} + + +} // namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_pcg_inner.hh b/RandLAPACK/drivers/rl_pcg_inner.hh new file mode 100644 index 000000000..a14d3632b --- /dev/null +++ b/RandLAPACK/drivers/rl_pcg_inner.hh @@ -0,0 +1,218 @@ +#pragma once + +// Internal component: pcg_inner, the instrumented conjugate-gradient kernel shared +// by IterRefineLSQ (rl_iter_refine_lsq.hh) and restarted_pcg_ne +// (rl_restarted_pcg_ne.hh), so the two least-squares drivers run the SAME inner +// solver: stagnation window with best-iterate return (from ISAAC diagnostic +// evidence), warm-start entry against the true residual, and per-solve +// diagnosis reporting. +// +// The kernel solves the SPD system M z = c on R^n given only a matvec callable +// apply_M(v, out). Tolerances are relative to ||c||; the caller decides what M, c, +// and the tolerance mean (IterRefineLSQ: correction equation at fixed inner_tol; +// restarted_pcg_ne: correction equation at the loose per-round drop). +// +// Relationship to RandLAPACK::pcg (comps/rl_determiter.hh): that solver is a +// general block/multi-RHS PCG over a caller-supplied preconditioner operator N +// and seminorm stopping test, with no stagnation window or best-iterate return. +// This kernel is a deliberately separate, narrower single-RHS CG built for the +// Q-less least-squares solvers: stagnation-window exit, best-iterate return, +// and a per-solve PCGInnerReport. Not a duplicate to consolidate. + +#include "rl_blaspp.hh" + +#include +#include +#include +#include + + +namespace RandLAPACK { + + +/// Exit condition of one inner-CG solve. Recorded per step/round so a capped, +/// non-converged solve is distinguishable from a converged one. +enum class InnerCGStatus : int { + Converged = 0, ///< reached the relative-residual target + HitCap = 1, ///< exhausted max_iters without reaching the target + Breakdown = 2, ///< p^T M p <= 0 (loss of orthogonality / non-SPD M) + Stagnated = 3 ///< residual stopped descending; exited early with the best iterate +}; + +/// What one inner-CG solve did, for diagnosis. +template +struct PCGInnerReport { + int iters = 0; + InnerCGStatus status = InnerCGStatus::Converged; + T relres = (T)0; ///< ||r||/||c|| at exit + T best_relres = (T)0; ///< smallest ||r||/||c|| seen + int best_iter = 0; ///< iteration achieving best_relres +}; + +/// Knobs of one inner-CG solve. Defaults match the values validated on the FEM2 +/// campaigns (stagnation window 20 at 0.1% improvement). +template +struct PCGInnerControls { + T tol; ///< stop when ||r|| <= tol * ||c|| + int max_iters; ///< hard iteration cap + int stag_window = 20; ///< <= 0 disables the stagnation exit + T stag_rel_improve = (T)1e-3; ///< drop counting as progress for the window + bool verbose = false; + const char* tag = "[PCG]"; ///< verbose-output prefix +}; + +/// @brief Instrumented CG on the SPD system M z = c. +/// +/// @param apply_M callable void(const T* v, T* out): out = M v. May use its own +/// scratch; must not alias v/out. +/// @param c right-hand side (length n). +/// @param n system dimension. +/// @param z solution buffer (length n). warm_start = false: initialized to +/// 0. warm_start = true: holds the starting iterate on entry; the +/// TRUE residual c - M z is computed (one extra M apply) and CG +/// runs from there, so a genuinely converged incoming iterate +/// returns immediately with 0 iterations. +/// @param cg_r,cg_p,cg_Mp,cg_zbest caller-allocated length-n workspaces. +/// @param ctl tolerances and caps (see PCGInnerControls). +/// @param rep filled with the solve's diagnosis (see PCGInnerReport). +/// @returns 0 on Converged/HitCap/Stagnated (the caller reads rep.status to tell +/// them apart); 1 on Breakdown. +/// +/// On Stagnated and HitCap exits z holds the BEST iterate seen, not the last one: +/// best_relres <= final relres by construction, and the 07-29 `bigcap` diagnostic +/// measured the last iterate to be up to 11x worse in outer solution error. +template +int pcg_inner(FApplyM&& apply_M, const T* c, int64_t n, + T* z, T* cg_r, T* cg_p, T* cg_Mp, T* cg_zbest, + const PCGInnerControls& ctl, PCGInnerReport& rep, + bool warm_start = false) +{ + T c_norm = blas::nrm2(n, c, 1); + T tol_abs = ctl.tol * c_norm; + if (c_norm == (T)0) { + // M is SPD, so M z = 0 has the unique solution z = 0. On a warm start the + // incoming z is already the previous attempt's answer to the same c = 0 + // system, i.e. 0, so writing 0 is correct on both paths. + std::fill(z, z + n, (T)0); + rep.iters = 0; + rep.status = InnerCGStatus::Converged; + rep.relres = (T)0; rep.best_relres = (T)0; rep.best_iter = 0; + return 0; + } + + rep.best_iter = 0; + + if (!warm_start) { + // Initial guess z = 0; r = c - M*z = c. + std::fill(z, z + n, (T)0); + std::fill(cg_zbest, cg_zbest + n, (T)0); + std::copy(c, c + n, cg_r); + rep.best_relres = (T)1; + } else { + // TRUE residual at the incoming iterate: r = c - M z. The best-iterate + // snapshot starts at z itself, so a restart can never end worse than + // where it began. + apply_M(z, cg_Mp); + for (int64_t i = 0; i < n; ++i) cg_r[i] = c[i] - cg_Mp[i]; + std::copy(z, z + n, cg_zbest); + T r0 = blas::nrm2(n, cg_r, 1); + rep.best_relres = r0 / c_norm; + if (r0 <= tol_abs) { + rep.iters = 0; + rep.status = InnerCGStatus::Converged; + rep.relres = rep.best_relres; + return 0; + } + } + std::copy(cg_r, cg_r + n, cg_p); + + // Stagnation state: `stag_ref` is the residual at the last SIGNIFICANT + // improvement (a drop of at least stag_rel_improve), and `last_improve_it` when + // it happened. A merely-noisy decrease does not count as progress: the + // pathological case descends by ~0 for hundreds of iterations. + T stag_ref = std::numeric_limits::max(); + int last_improve_it = 0; + + T rs_old = blas::dot(n, cg_r, 1, cg_r, 1); + + for (int it = 0; it < ctl.max_iters; ++it) { + apply_M(cg_p, cg_Mp); + + T pMp = blas::dot(n, cg_p, 1, cg_Mp, 1); + if (!(pMp > 0)) { + // Hand back the BEST iterate here too, not the last one, for + // consistency with the Stagnated/HitCap exits. rep.relres matches + // the returned iterate. + // Note the M apply of this aborted iteration ran but is not counted + // in rep.iters (iters = COMPLETED CG iterations, everywhere). + std::copy(cg_zbest, cg_zbest + n, z); + rep.iters = it; + rep.status = InnerCGStatus::Breakdown; + rep.relres = rep.best_relres; + return 1; // CG breakdown (loss of orthogonality / non-SPD M) + } + T alpha = rs_old / pMp; + + blas::axpy(n, alpha, cg_p, 1, z, 1); // z ← z + alpha p + blas::axpy(n, -alpha, cg_Mp, 1, cg_r, 1); // r ← r - alpha Mp + + T rs_new = blas::dot(n, cg_r, 1, cg_r, 1); + T r_norm = std::sqrt(rs_new); + T relres = r_norm / c_norm; + if (relres < rep.best_relres) { + rep.best_relres = relres; + rep.best_iter = it + 1; + std::copy(z, z + n, cg_zbest); // snapshot for the stagnation exit + } + if (relres < stag_ref * ((T)1 - ctl.stag_rel_improve)) { + stag_ref = relres; + last_improve_it = it + 1; + } + + if (ctl.verbose) { + std::printf("%s inner CG iter %d: ||r||/||c|| = %.4e\n", + ctl.tag, it + 1, (double)relres); + } + // Convergence is checked BEFORE stagnation: a solve that reaches the target + // reports Converged even if its last few steps were flat. + if (r_norm <= tol_abs) { + rep.iters = it + 1; + rep.status = InnerCGStatus::Converged; + rep.relres = relres; + return 0; + } + if (ctl.stag_window > 0 && + (it + 1) - last_improve_it >= ctl.stag_window) { + // Residual has flatlined. More iterations cannot reach the target, and + // are measurably harmful to the outer solution, so stop and hand back + // the best iterate rather than the last one. + std::copy(cg_zbest, cg_zbest + n, z); + rep.iters = it + 1; + rep.status = InnerCGStatus::Stagnated; + rep.relres = rep.best_relres; + if (ctl.verbose) { + std::printf("%s inner CG STAGNATED at iter %d " + "(no %.1e improvement in %d iters); returning best " + "iterate from iter %d, relres %.4e\n", + ctl.tag, it + 1, (double)ctl.stag_rel_improve, + ctl.stag_window, rep.best_iter, (double)rep.best_relres); + } + return 0; + } + + T beta = rs_new / rs_old; + for (int64_t i = 0; i < n; ++i) cg_p[i] = cg_r[i] + beta * cg_p[i]; + rs_old = rs_new; + } + // Exhausted the budget without reaching the target. Still returns 0, because a + // capped solve is not necessarily an error for the caller; the status records + // it. Hand back the best iterate here too. + std::copy(cg_zbest, cg_zbest + n, z); + rep.iters = ctl.max_iters; + rep.status = InnerCGStatus::HitCap; + rep.relres = rep.best_relres; + return 0; +} + + +} // namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_restarted_pcg_ne.hh b/RandLAPACK/drivers/rl_restarted_pcg_ne.hh new file mode 100644 index 000000000..6e426e2d4 --- /dev/null +++ b/RandLAPACK/drivers/rl_restarted_pcg_ne.hh @@ -0,0 +1,452 @@ +#pragma once + +// Public API: restarted_pcg_ne, restarted PCG on the right-preconditioned normal +// equations, with an optional upper-triangular RIGHT preconditioner R. +// +// Solves min_x ||b - A x||_2 for a tall LinearOperator A (m x n, m >= n) by +// running conjugate gradients on the preconditioned normal equations +// +// H z = g, H = R^{-T} A^T A R^{-1}, g = R^{-T} A^T b, +// +// and recovering x = R^{-1} z. Pass R = nullptr for the unpreconditioned normal +// equations (H = A^T A). This is the second solver of the reference Toeplitz +// least-squares benchmark (ar_sysid_toeplitz_qless_qr_benchmark.m, +// restarted_pcg_preconditioned_normal_eq), ported with the same semantics: +// +// * Each outer round runs CG on the correction equation H dz = r_ne with a +// deliberately LOOSE relative tolerance `restart_drop` (default 1e-4, i.e. +// stop after a 1e4x residual drop) and an inner cap of +// min(restart_maxit, max_iters - iters_so_far). The inner solver is the +// shared instrumented kernel (rl_pcg_inner.hh), the same CG that +// IterRefineLSQ runs, with stagnation window and best-iterate return. +// That is a deliberate deviation from the reference's plain MATLAB pcg: a +// round whose target is unreachable exits at its residual floor with the +// best iterate instead of grinding out the full cap, and stagnation is not +// treated as terminal (the next round's true-residual restart decides). +// * After the round, z += dz and BOTH residuals are recomputed exactly: the +// true least-squares residual b - A x, and the normal-equation residual in +// the STABLE form R^{-T}(A^T(b - A x)) (Epperly et al. Alg. 1 line 5) rather +// than the reference's g - H z. The two are mathematically identical, but +// g - H z subtracts large kappa-contaminated quantities and floors the +// achievable accuracy on hard problems (second deliberate deviation from the +// reference; measured A/B in the round-residual comment below). +// * The loop exits when ||b - A x|| / ||b|| <= tol (success), when the TOTAL +// inner-iteration budget max_iters is exhausted, when the round budget +// max_restarts is spent, when the inner CG breaks down (indefinite H apply, +// a sign the factor R is unusable), or when outer_stag_window consecutive +// rounds fail to improve the true LS residual (the LS-floor exit). Each exit +// has its own status code; see @returns. +// +// The convergence test is on the true LS residual, matching the reference; the +// inner drop tolerance only paces the restarts. + +#include "rl_blaspp.hh" +#include "rl_blas2_threads.hh" +#include "rl_exceptions.hh" +#include "rl_pcg_inner.hh" +#include "../linops/rl_concepts.hh" + +#include +#include +#include +#include +#include + + +namespace RandLAPACK { + + +/// Per-round records of one restarted_pcg_ne run, for callers that need more +/// than the aggregate outputs (IterRefineLSQ delegates here and republishes +/// these as its per-step diagnostics). All vectors have one entry per round. +/// The t_* fields separate work done INSIDE the inner CG kernel from the +/// restart loop's own residual recomputations, so a non-overlapping timing +/// breakdown can be assembled by the caller. +template +struct PCGRoundHistory { + std::vector iters; ///< inner CG iterations of the round + std::vector status; ///< InnerCGStatus of the round (as int) + std::vector relres; ///< kernel relres of the RETURNED iterate (the + ///< best-iterate relres on Stagnated/HitCap exits) + std::vector best_relres; ///< best kernel relres seen in the round + std::vector best_iter; ///< iteration achieving best_relres + std::vector ls_relres; ///< true LS relres after the round + long t_inner_us = 0; ///< wallclock inside pcg_inner + long t_fwd_inner_us = 0; ///< A applies inside the kernel + long t_adj_inner_us = 0; ///< A^T applies inside the kernel + long t_trsm_inner_us = 0; ///< trsv time inside the kernel + void clear() { + iters.clear(); status.clear(); relres.clear(); + best_relres.clear(); best_iter.clear(); ls_relres.clear(); + t_inner_us = t_fwd_inner_us = t_adj_inner_us = t_trsm_inner_us = 0; + } +}; + + +/// @brief Restarted PCG on the right-preconditioned normal equations for +/// min ||b - A x||, optional upper-triangular right preconditioner R. +/// +/// @param[in] A tall LinearOperator (m x n), applied as A*v and A^T*u. +/// @param[in] m,n dimensions (m >= n). +/// @param[in] R upper-triangular right preconditioner (n x n, ColMajor) +/// or nullptr for none. Must be nonsingular when supplied. +/// @param[in] ldr leading dimension of R. +/// @param[in] b right-hand side (length m). +/// @param[out] x solution (length n). +/// @param[in] tol target on the true LS relative residual ||b - A x|| / ||b||. +/// @param[in] max_iters TOTAL inner CG iteration budget, shared across restarts. +/// @param[out] iters_done total inner CG iterations actually run. +/// @param[in] restart_maxit inner CG cap per restart (reference default 200). +/// @param[in] restart_drop inner CG relative residual drop per restart, in (0,1). +/// @param[in] max_restarts additional outer rounds allowed after the first, the +/// IterRefineLSQ inner_restarts convention: 0 means a single +/// round, 3 means up to four rounds, negative means unlimited +/// (the reference behaviour, rounds bounded only by max_iters). +/// @param[out] restarts_done optional: number of outer restart rounds (may be nullptr). +/// @param[out] times optional [fwd_us, adj_us, trsm_us, total_us] (may be nullptr). +/// @param[out] final_relres optional: the true LS relative residual at termination. +/// @param[in] stag_window,stag_rel_improve stagnation exit knobs, forwarded to +/// the inner kernel (see PCGInnerControls). +/// @param[in] inner_abs_tol ABSOLUTE inner target, relative to the INITIAL +/// normal-equation right-hand side ||g||. When > 0, a +/// round whose NE residual has already fallen to +/// inner_abs_tol * ||g|| stops immediately instead of +/// grinding for a further restart_drop factor (the "CG +/// still terminates once below the target" guard). 0 +/// disables the guard. +/// @param[out] history optional per-round records (see PCGRoundHistory). +/// @param[in] x0 optional initial guess (length n). nullptr = cold start +/// (x = 0), the historical behaviour and the policy for every +/// Q-less method. Supplied, the solver refines THAT iterate: +/// z is seeded with R x0 so that x = R^{-1} z reproduces it, +/// and the first round's normal-equation residual is taken +/// from the true residual b - A x0 rather than from g. +/// Lets a solver's own answer (e.g. Blendenpik's) be +/// handed to iterative refinement, separating +/// preconditioner quality from solver structure. +/// @param[in] outer_stag_window consecutive rounds without a stag_rel_improve +/// drop of the TRUE LS residual (measured against the last +/// significant improvement, mirroring the inner kernel) that +/// end the loop as an LS-floor exit. <= 0 disables the outer +/// exit. Decoupled from stag_window: the two mechanisms +/// are independent. +/// @returns 0 if the LS tolerance was met; +/// 1 if the total inner-iteration budget was exhausted; +/// 2 if the inner CG broke down or made no progress (reference flag 2); +/// 3 if the outer round budget (max_restarts) was spent; +/// 4 if the run ended at its LS floor (outer stagnation exit, or an +/// exactly-zero NE residual with the LS tolerance still unmet). +/// Codes 3 and 4 are distinct from 1; callers that only test +/// zero/nonzero are unaffected. +template +int restarted_pcg_ne( + GLO& A, int64_t m, int64_t n, + const T* R, int64_t ldr, + const T* b, T* x, + T tol, int max_iters, + int& iters_done, + int restart_maxit = 200, + T restart_drop = (T)1e-4, + int max_restarts = -1, + int* restarts_done = nullptr, + long* times = nullptr, + T* final_relres = nullptr, + int stag_window = 20, + T stag_rel_improve = (T)1e-3, + T inner_abs_tol = (T)0, + PCGRoundHistory* history = nullptr, + const T* x0 = nullptr, + int outer_stag_window = 2) +{ + randlapack_require(restart_drop > (T)0 && restart_drop < (T)1) + << "restarted_pcg_ne: restart_drop must lie in (0,1)"; + randlapack_require(restart_maxit >= 1) + << "restarted_pcg_ne: restart_maxit must be >= 1"; + + using clock = std::chrono::steady_clock; + using std::chrono::duration_cast; + using std::chrono::microseconds; + // One width for the whole solve: narrows width-capped kernels (the Toeplitz + // FFT) to the trsv width so the inner loop pays no team re-formation per + // width transition. No-op for operators that do not consult the context. + SolveWidthScope solve_scope(n); + long t_fwd = 0, t_adj = 0, t_trsm = 0; + // Inner/outer attribution: apply_H is called both inside the CG kernel and + // by the restart loop's residual recomputations. The flag routes each op's + // time into the inner-only counters too, so history (when requested) can + // report a non-overlapping breakdown. + bool in_kernel = false; + long t_fwd_in = 0, t_adj_in = 0, t_trsm_in = 0, t_kernel = 0; + auto total_start = clock::now(); + if (history) history->clear(); + + const bool prec = (R != nullptr); + + // Workspaces (raw T*, freed on every return path via cleanup()). + T* z = new T[n](); // preconditioned solution, x = R^{-1} z + T* g = new T[n](); // normal-equation right-hand side R^{-T} A^T b + T* r_ne = new T[n](); // NE residual in the stable form R^{-T} A^T (b - A x) + T* dz = new T[n](); // inner CG correction + T* p = new T[n](); // CG direction + T* q = new T[n](); // holds H p + T* r = new T[n](); // inner CG (recursive) residual + T* zb = new T[n](); // kernel best-iterate snapshot + T* sc = new T[n](); // trsv scratch + T* wm = new T[m](); // length-m scratch (A applies) + auto cleanup = [&]() { delete[] z; delete[] g; delete[] r_ne; delete[] dz; + delete[] p; delete[] q; delete[] r; delete[] zb; + delete[] sc; delete[] wm; }; + + // H v = R^{-T} (A^T (A (R^{-1} v))) (out has length n; out may not alias v) + auto apply_H = [&](const T* vin, T* out) { + long dt; + const T* fwd_in = vin; + if (prec) { + std::copy(vin, vin + n, sc); // sc = v + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); // cap threads: see rl_blas2_threads.hh + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, n, R, ldr, sc, 1); + } + dt = duration_cast(clock::now() - ts).count(); + t_trsm += dt; if (in_kernel) t_trsm_in += dt; + fwd_in = sc; // sc = R^{-1} v + } + auto tf = clock::now(); + A(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, fwd_in, n, (T)0.0, wm, m); + dt = duration_cast(clock::now() - tf).count(); + t_fwd += dt; if (in_kernel) t_fwd_in += dt; + auto ta = clock::now(); + A(blas::Side::Left, blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + n, 1, m, (T)1.0, wm, m, (T)0.0, out, n); + dt = duration_cast(clock::now() - ta).count(); + t_adj += dt; if (in_kernel) t_adj_in += dt; + if (prec) { + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); // cap threads: see rl_blas2_threads.hh + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::Trans, blas::Diag::NonUnit, n, R, ldr, out, 1); + } + dt = duration_cast(clock::now() - ts).count(); + t_trsm += dt; if (in_kernel) t_trsm_in += dt; + } + }; + + // True LS relative residual ||b - A x|| / ||b|| for the CURRENT z (recovers x too). + T bnorm = blas::nrm2(m, b, 1); + T bden = std::max(bnorm, std::numeric_limits::min()); + auto recover_x_and_relres = [&]() -> T { + std::copy(z, z + n, x); + if (prec) { + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); // cap threads: see rl_blas2_threads.hh + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, n, R, ldr, x, 1); + } + t_trsm += duration_cast(clock::now() - ts).count(); + } + auto tf = clock::now(); + A(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, x, n, (T)0.0, wm, m); + t_fwd += duration_cast(clock::now() - tf).count(); + blas::scal(m, (T)-1.0, wm, 1); + blas::axpy(m, (T)1.0, b, 1, wm, 1); // wm = b - A x + return blas::nrm2(m, wm, 1) / bden; + }; + + // g = R^{-T} A^T b. + { + auto ta = clock::now(); + A(blas::Side::Left, blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + n, 1, m, (T)1.0, b, m, (T)0.0, g, n); + t_adj += duration_cast(clock::now() - ta).count(); + if (prec) { + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); // cap threads: see rl_blas2_threads.hh + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::Trans, blas::Diag::NonUnit, n, R, ldr, g, 1); + } + t_trsm += duration_cast(clock::now() - ts).count(); + } + } + + // Cold start: z = 0, so x = 0 and the NE residual is exactly g. + // Warm start: seed z = R x0 so that x = R^{-1} z recovers x0, then take the + // NE residual from the TRUE residual b - A x0 in the same stable form the + // restart loop uses (g - H z would reintroduce the cancellation this + // stable form avoids). + T relres; + if (x0 == nullptr) { + // z = 0, so x = 0 exactly and ||b - A x|| / ||b|| = 1 with no arithmetic: + // the full apply the recovery lambda would burn here (one FFT-class + // operator apply plus a trsv on a zero vector) is skipped. + std::copy(g, g + n, r_ne); + std::fill(x, x + n, (T)0); + relres = (bnorm > (T)0) ? (T)1 : (T)0; + } else { + std::copy(x0, x0 + n, z); + if (prec) { + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); // trmv degrades unguarded like trsv + blas::trmv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, n, R, ldr, z, 1); + } + t_trsm += duration_cast(clock::now() - ts).count(); + } + relres = recover_x_and_relres(); // leaves wm = b - A x + auto ta = clock::now(); + A(blas::Side::Left, blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + n, 1, m, (T)1.0, wm, m, (T)0.0, r_ne, n); + t_adj += duration_cast(clock::now() - ta).count(); + if (prec) { + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::Trans, blas::Diag::NonUnit, n, R, ldr, r_ne, 1); + } + t_trsm += duration_cast(clock::now() - ts).count(); + } + } + + // NE-space reference scale for the inner_abs_tol guard: the initial NE + // right-hand side ||g||. + T g0_norm = blas::nrm2(n, g, 1); + + iters_done = 0; + int restarts = 0; + int status = 1; + + // Outer stagnation state: when tol sits below the problem's achievable LS + // floor (e.g. a data noise floor above the requested tolerance), every + // round reaches the floor and further rounds burn iterations without + // progress. outer_stag_window consecutive rounds without a significant + // CUMULATIVE improvement over the last significant drop end the loop as + // an LS-floor exit (status 4). The reference advances only on a + // significant improvement, mirroring the inner kernel: a per-round + // reference update would kill steady slow descent (e.g. 0.05% per round) + // after two rounds. Disabled when outer_stag_window <= 0. + T ls_stag_ref = relres; + int ls_flat_rounds = 0; + + while (relres > tol && iters_done < max_iters) { + if (max_restarts >= 0 && restarts > max_restarts) { status = 3; break; } // round budget spent + T ne_norm = blas::nrm2(n, r_ne, 1); + if (ne_norm == (T)0) { + // Exactly-zero NE residual: x is the LS minimizer. That meets the + // caller's tolerance only if the LS residual itself does; otherwise + // this is the LS floor, not convergence. + status = (relres <= tol) ? 0 : 4; + break; + } + + ++restarts; + int inner_cap = std::min(restart_maxit, max_iters - iters_done); + + // ---- Inner CG on H dz = r_ne, from dz = 0 (shared instrumented kernel): + // loose target restart_drop * ||r_ne||, stagnation window + best- + // iterate return. Breakdown maps to the reference's terminal flag 2; + // Stagnated/HitCap continue to the next true-residual round. ---- + PCGInnerControls ctl; + ctl.tol = restart_drop; + // Absolute-target guard: once ||r_ne|| has fallen to inner_abs_tol * ||g||, + // the round's effective target is already met (or nearly so) and grinding + // out a further restart_drop factor is wasted work at the noise floor. The + // kernel tolerance is relative to THIS round's RHS, so rescale. + // ne_norm > 0 is not re-checked here: the loop already broke above on + // ne_norm == 0, so every reach of this point has ne_norm > 0. + if (inner_abs_tol > (T)0) { + T floor_rel = inner_abs_tol * g0_norm / ne_norm; + if (floor_rel > ctl.tol) ctl.tol = floor_rel; + } + ctl.max_iters = inner_cap; + ctl.stag_window = stag_window; + ctl.stag_rel_improve = stag_rel_improve; + ctl.tag = "[PCG-NE]"; + PCGInnerReport rep; + in_kernel = true; + auto tk0 = clock::now(); + int kret = pcg_inner(apply_H, r_ne, n, dz, r, p, q, zb, ctl, rep, + /*warm_start=*/false); + t_kernel += duration_cast(clock::now() - tk0).count(); + in_kernel = false; + int inner_iters = rep.iters; + + blas::axpy(n, (T)1.0, dz, 1, z, 1); + iters_done += inner_iters; + + // Recompute BOTH residuals exactly; this is the restart that removes + // recursive-residual drift (the point of the algorithm). + relres = recover_x_and_relres(); + // STABLE residual form: map wm = b - A x through R^{-T} A^T (Epperly + // et al. Alg. 1 line 5) rather than the reference's g - H z, which + // subtracts two large kappa-contaminated quantities and floors the + // achievable accuracy on hard problems. + { + auto ta = clock::now(); + A(blas::Side::Left, blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + n, 1, m, (T)1.0, wm, m, (T)0.0, r_ne, n); + t_adj += duration_cast(clock::now() - ta).count(); + if (prec) { + auto ts = clock::now(); + { Blas2ThreadGuard tg(n); // cap threads: see rl_blas2_threads.hh + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::Trans, blas::Diag::NonUnit, n, R, ldr, r_ne, 1); + } + t_trsm += duration_cast(clock::now() - ts).count(); + } + } + + if (history) { + history->iters.push_back(inner_iters); + history->status.push_back(static_cast(rep.status)); + history->relres.push_back(rep.relres); + history->best_relres.push_back(rep.best_relres); + history->best_iter.push_back(rep.best_iter); + history->ls_relres.push_back(relres); + } + + if (relres <= tol) { status = 0; break; } + if (kret != 0) { status = 2; break; } // breakdown: R unusable + // A zero-iteration round changes nothing, so the loop must end either + // way; gate the MEANING on the kernel status, not the count: a round + // whose target was already met at entry is the LS floor (the NE + // target cannot improve the true residual further), not a breakdown. + // Defensive: unreachable today (warm_start is hard-wired false above and + // ne_norm > 0 is guaranteed by the break earlier in this loop, so + // pcg_inner cannot return 0 iterations here); guards a future kernel + // call that warm-starts dz and could legitimately return 0 iterations. + if (inner_iters == 0) { + status = (rep.status == InnerCGStatus::Converged) ? 4 : 2; + break; + } + if (outer_stag_window > 0) { + if (relres >= ls_stag_ref * ((T)1 - stag_rel_improve)) { + if (++ls_flat_rounds >= outer_stag_window) { + status = 4; break; // LS floor reached: stop + } + } else { + ls_flat_rounds = 0; + ls_stag_ref = relres; // reference advances on significant drops only + } + } + } + + if (relres <= tol) status = 0; + + if (restarts_done) *restarts_done = restarts; + if (times) { times[0] = t_fwd; times[1] = t_adj; times[2] = t_trsm; + times[3] = duration_cast(clock::now() - total_start).count(); } + if (history) { + history->t_inner_us = t_kernel; + history->t_fwd_inner_us = t_fwd_in; + history->t_adj_inner_us = t_adj_in; + history->t_trsm_inner_us = t_trsm_in; + } + if (final_relres) *final_relres = relres; + cleanup(); + return status; +} + + +} // namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_scholqr3_linops.hh b/RandLAPACK/drivers/rl_scholqr3_linops.hh index abbef9ad5..c553d1bf7 100644 --- a/RandLAPACK/drivers/rl_scholqr3_linops.hh +++ b/RandLAPACK/drivers/rl_scholqr3_linops.hh @@ -4,6 +4,7 @@ #include "rl_blaspp.hh" #include "rl_lapackpp.hh" #include "rl_linops.hh" +#include "../comps/rl_cholqr.hh" #include #include @@ -11,38 +12,26 @@ #include #include #include +#include +#include + +// scholqr3_eps_shift() (the RANDLAPACK_SCHOLQR3_SHIFT env knob) now lives +// in comps/rl_cholqr.hh, shared with the dense sCholQR3 driver. using namespace std::chrono; namespace RandLAPACK { -/// Shifted Cholesky QR3 for abstract linear operators. -/// -/// Linop analogue of shifted CholQR3: computes A = QR where A is any type -/// satisfying the LinearOperator concept. All Gram matrices are formed through -/// A(NoTrans, ...) and A(Trans, ...) calls, so A can be dense, sparse, -/// composite, or any other operator type without modification. -/// -/// Fully-blocked implementation: never materializes the full m x n operator product -/// during the QR iterations. All three iterations compute their Gram matrices through -/// blocked linop calls, using an accumulated right-factor M = R1^{-1} R2^{-1} ... -/// to avoid storing Q explicitly. -/// -/// Algorithm: -/// 1. Shifted CholQR1: G1 = A^T A + s*I, R1 = chol(G1), M <- R1^{-1} -/// 2. CholQR2: G2 = M^T A^T A M, R2 = chol(G2), R = R2*R1, M <- M*R2^{-1} -/// 3. CholQR3: G3 = M^T A^T A M, R3 = chol(G3), R = R3*R +/// Shifted Cholesky QR3 for abstract linear operators (fully-blocked variant). /// -/// Blocked Gram computation for iteration k (M_k = accumulated R-inverse): -/// for each column block j of width b: -/// W = A * M_k[:, j:j+b] (linop NoTrans, m x b) -/// Z = A^T * W (linop Trans, n x b) -/// G[:, j:j+b] = M_k^T * Z (gemm, n x b) +/// Algorithm 3 from the collaborator's spec: +/// iter 1: cholqr_primitive(A) with shift s = eps * ||A||_F^2 -> R_1 +/// iter i = 2, 3: cholqr_primitive(A, R_{i-1}, TRSM_IDENTITY) -> R_i +/// return R_3 /// -/// Peak memory: O(m*b + n^2) -- no m x n buffer needed during QR iterations. -/// If test_mode is enabled, Q = A * R^{-1} is materialized at the end (m x n). -/// -/// The shift is computed as: shift = 11 * eps * n * ||A||_F^2 +/// Peak memory O(n^2 + (m+n)*b_eff): never materializes the m × n operator product +/// during the QR iterations. If test_mode is enabled, Q = A * R^{-1} is materialized +/// at the end (m × n, outside the timing region). /// /// Reference: Shifted Cholesky QR from Fukaya et al. (SISC, 2020). /// @@ -52,48 +41,60 @@ class sCholQR3_linops { bool timing; bool test_mode; - T eps; // Q-factor for test mode (only allocated if test_mode = true) T* Q; int64_t Q_rows; int64_t Q_cols; - // Individual Cholesky factors from each iteration (n x n upper triangular). - std::vector G1_factor; - std::vector G2_factor; - std::vector G3_factor; - - // Timing breakdown (18 entries): - // [0] alloc - buffer allocation - // [1] fwd1 - Iter 1 NoTrans: A * M[:, block] - // [2] adj1 - Iter 1 Trans: A^T * W (direct to G since M=I) - // [3] chol1 - Iter 1 Cholesky (potrf) - // [4] upd1 - M = R1^{-1} (n x n trsm) - // [5] fwd2 - Iter 2 NoTrans: A * M[:, block] - // [6] adj2 - Iter 2 Trans: A^T * W - // [7] gemm2 - Iter 2 M^T * Z - // [8] chol2 - Iter 2 Cholesky - // [9] upd2 - R = R2*R1, M *= R2^{-1} - // [10] fwd3 - Iter 3 NoTrans - // [11] adj3 - Iter 3 Trans - // [12] gemm3 - Iter 3 M^T * Z - // [13] chol3 - Iter 3 Cholesky - // [14] upd3 - R = R3*R - // [15] q_mat - Q materialization for test mode (0 if not test_mode) - // [16] rest - unaccounted time - // [17] total - wall-clock total + // Timing breakdown (18 entries; layout preserved for matlab plotters): + // [0] alloc + // [1] fwd1 [2] adj1 [3] chol1 [4] upd1 + // [5] fwd2 [6] adj2 [7] gemm2 [8] chol2 [9] upd2 + // [10] fwd3 [11] adj3 [12] gemm3 [13] chol3 [14] upd3 + // [15] q_mat [16] rest [17] total std::vector times; + /// Total measured wall-clock (microseconds) of the last call(), or -1 if timing + /// was off. Every driver in this family packs the total as the LAST times[] entry, + /// but the entry COUNT differs per driver (6 / 11 / 15 / 18). Callers used to hard- + /// code that index (times[5], times[10], times[14], times[17]), so adding or + /// removing one slot silently wrote the wrong number into every CSV with no compile + /// error. Read the total through here instead. + long total_us() const { return times.empty() ? -1L : times.back(); } + + int64_t block_size; - // Column-block size for blocked Gram computations. + // Adaptive-shift policy (see cholqr_primitive). Shift s = factor * trace(G). // - // Controls the width of column blocks used in all three iterations' - // Gram matrix computations. Smaller values reduce peak memory - // (O(m*block_size + n^2) instead of O(m*n)), at the cost of - // more linop calls (2 * ceil(n/block_size) per iteration). + // iter 1: shifted. shift_factor_iter1 < 0 (the default) resolves at call + // time to the paper's prescription 11 * n * eps (FukayaEtAl2020, c = 11), + // or to plain eps when RANDLAPACK_SCHOLQR3_SHIFT=eps is set. The eps + // variant is a smaller shift kept for A/B campaigns: a shift far above + // σ_min²(A) makes the iter-2 Gram rank-deficient in principle, though + // with unshifted refinement passes plus the adaptive retry the paper + // shift has measured clean on the FEM2 campaigns. A caller-set value + // >= 0 is used verbatim. // - // When block_size <= 0 or >= n, uses b_eff = n (single block per loop). - int64_t block_size; + // iters 2-3: UNSHIFTED (shift_factor_iter23 = 0). This is the defining + // feature of Fukaya shifted-CholeskyQR3: the refinement passes are plain + // CholeskyQR2, which is what drives orthogonality down to machine level. + // A persistent eps shift on these passes (the old setting) never gets + // removed: it floors orth at ~2n*eps (≈1.8e-12 in double) and, in single, + // over-regularizes R into a useless preconditioner (CG stalls, 100s of + // inner iters). The adaptive retry below still shifts a refinement pass + // *only* if its potrf genuinely fails. + // + // The retry loop bumps shift × 10 if potrf bails, unboundedly + // (max_retries = -1) until the Gram is PD. + T shift_factor_iter1; + T shift_factor_iter23; + int max_retries; + T shift_growth; + int n_chol_retries = 0; ///< shift retries used on the last call (0 = clean) + /// Per-pass shift record from the last call (passes 1-3): absolute shift the + /// successful potrf carried (0 = unshifted) and that pass's Gram trace. + T chol_applied_shifts[3] = {T(0), T(0), T(0)}; + T chol_gram_traces[3] = {T(0), T(0), T(0)}; sCholQR3_linops( bool time_subroutines, @@ -101,12 +102,16 @@ class sCholQR3_linops { bool enable_test_mode = false ) { timing = time_subroutines; - eps = ep; - block_size = 0; + (void)ep; // kept in the signature for call-site compatibility; unused + block_size = kDefaultGramBlockSize; test_mode = enable_test_mode; Q = nullptr; Q_rows = 0; Q_cols = 0; + shift_factor_iter1 = T(-1); // < 0: resolve default (11*n*eps, or eps via env) at call time + shift_factor_iter23 = T(0); // iters 2-3 unshifted (Fukaya); retry covers genuine non-PD + max_retries = -1; // unbounded retries (no ceiling), consistent with CholQR/CholQR2 + shift_growth = T(10); } ~sCholQR3_linops() { @@ -115,351 +120,86 @@ class sCholQR3_linops { } } - /// Computes the QR factorization A = QR using shifted Cholesky QR3. - /// - /// @param[in] A - /// The m-by-n linear operator (m and n read from A.n_rows, A.n_cols). - /// - /// @param[out] R - /// Pre-allocated n-by-n buffer. On exit, stores the upper-triangular - /// R factor. Zero entries are not compressed. - /// - /// @param[in] ldr - /// Leading dimension of R. - /// - /// @return = 0: successful exit template int call( GLO& A, T* R, int64_t ldr ) { - ///--------------------TIMING VARS--------------------/ - steady_clock::time_point t_start, t_stop; - steady_clock::time_point total_t_start, total_t_stop; - long alloc_dur = 0; - long fwd1_dur = 0, adj1_dur = 0, chol1_dur = 0, upd1_dur = 0; - long fwd2_dur = 0, adj2_dur = 0, gemm2_dur = 0, chol2_dur = 0, upd2_dur = 0; - long fwd3_dur = 0, adj3_dur = 0, gemm3_dur = 0, chol3_dur = 0, upd3_dur = 0; - long q_mat_dur = 0, total_dur = 0; - - if(this->timing) - total_t_start = steady_clock::now(); + steady_clock::time_point t0, t1, total_t_start, total_t_stop; + long q_mat_dur = 0; + + if (this->timing) total_t_start = steady_clock::now(); int64_t m = A.n_rows; int64_t n = A.n_cols; - - // Determine effective block width. int64_t b_eff = (this->block_size > 0 && this->block_size < n) ? this->block_size : n; - if(this->timing) - t_start = steady_clock::now(); - - // ---- Allocate buffers ---- - // G: n x n Gram matrix / Cholesky workspace (zero-init for lower triangle) - T* G = new T[n * n](); - - // R_temp: n x n workspace for R accumulation via trmm - T* R_temp = new T[n * n](); - - // M: n x n accumulated R-inverse product (starts as identity) - T* M = new T[n * n](); - RandLAPACK::util::eye(n, n, M); - - // A_temp: m x b_eff buffer for linop NoTrans output - T* A_temp = new T[m * b_eff]; - - // Z_buf: n x b_eff buffer for linop Trans output (used in iterations 2-3) - T* Z_buf = new T[n * b_eff]; - - if(this->timing) { - t_stop = steady_clock::now(); - alloc_dur = duration_cast(t_stop - t_start).count(); - } - - //================================================================ - // Iteration 1: Shifted Cholesky QR - //================================================================ - // Blocked Gram: G = A^T A (since M = I, no M^T multiply needed) - long fwd1_accum = 0, adj1_accum = 0; - for (int64_t j = 0; j < n; j += b_eff) { - int64_t b_j = std::min(b_eff, n - j); - - // W = A * M[:, j:j+b] (= A * I[:, j:j+b] since M = I) - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, b_j, n, (T)1.0, M + j * n, n, (T)0.0, A_temp, m); - if(this->timing) { t_stop = steady_clock::now(); fwd1_accum += duration_cast(t_stop - t_start).count(); } - - // G[:, j:j+b] = A^T * W (direct to G since M = I) - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, - n, b_j, m, (T)1.0, A_temp, m, (T)0.0, G + j * n, n); - if(this->timing) { t_stop = steady_clock::now(); adj1_accum += duration_cast(t_stop - t_start).count(); } - } - - // Compute shift from ||A||_F^2 = trace(G) - T norm_A_sq = 0; - for (int64_t i = 0; i < n; ++i) - norm_A_sq += G[i * (n + 1)]; - T shift = 11 * std::numeric_limits::epsilon() * n * norm_A_sq; - - // Add shift to diagonal: G = G + shift * I - for (int64_t i = 0; i < n; ++i) - G[i * (n + 1)] += shift; - - if(this->timing) { - fwd1_dur = fwd1_accum; - adj1_dur = adj1_accum; - t_start = steady_clock::now(); - } - - // Zero lower triangle, Cholesky: G = R1^T * R1 - if (n > 1) - lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &G[1], n); - if (lapack::potrf(Uplo::Upper, n, G, n)) { - delete[] G; delete[] R_temp; delete[] M; - delete[] A_temp; delete[] Z_buf; - return 1; - } - - // Save G1 factor - this->G1_factor.resize(n * n, (T)0.0); - lapack::lacpy(MatrixType::Upper, n, n, G, n, this->G1_factor.data(), n); - - // Initialize R = R1 - lapack::lacpy(MatrixType::Upper, n, n, G, n, R, ldr); - - if(this->timing) { - t_stop = steady_clock::now(); - chol1_dur = duration_cast(t_stop - t_start).count(); - t_start = steady_clock::now(); - } - - // Update M: M = I * R1^{-1} = R1^{-1} - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, n, n, (T)1.0, G, n, M, n); - - if(this->timing) { - t_stop = steady_clock::now(); - upd1_dur = duration_cast(t_stop - t_start).count(); - } - - //================================================================ - // Iteration 2: Cholesky QR - //================================================================ - // Blocked Gram: G2 = M^T * A^T * A * M where M = R1^{-1} - long fwd2_accum = 0, adj2_accum = 0, gemm2_accum = 0; - for (int64_t j = 0; j < n; j += b_eff) { - int64_t b_j = std::min(b_eff, n - j); - - // W = A * M[:, j:j+b] - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, b_j, n, (T)1.0, M + j * n, n, (T)0.0, A_temp, m); - if(this->timing) { t_stop = steady_clock::now(); fwd2_accum += duration_cast(t_stop - t_start).count(); } - - // Z = A^T * W - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, - n, b_j, m, (T)1.0, A_temp, m, (T)0.0, Z_buf, n); - if(this->timing) { t_stop = steady_clock::now(); adj2_accum += duration_cast(t_stop - t_start).count(); } - - // G[:, j:j+b] = M^T * Z - if(this->timing) t_start = steady_clock::now(); - blas::gemm(Layout::ColMajor, Op::Trans, Op::NoTrans, - n, b_j, n, (T)1.0, M, n, Z_buf, n, (T)0.0, G + j * n, n); - if(this->timing) { t_stop = steady_clock::now(); gemm2_accum += duration_cast(t_stop - t_start).count(); } - } - - if(this->timing) { - fwd2_dur = fwd2_accum; - adj2_dur = adj2_accum; - gemm2_dur = gemm2_accum; - t_start = steady_clock::now(); - } - - // Zero lower triangle, Cholesky: G = R2^T * R2 - if (n > 1) - lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &G[1], n); - if (lapack::potrf(Uplo::Upper, n, G, n)) { - delete[] G; delete[] R_temp; delete[] M; - delete[] A_temp; delete[] Z_buf; - return 2; - } - - // Save G2 factor - this->G2_factor.resize(n * n, (T)0.0); - lapack::lacpy(MatrixType::Upper, n, n, G, n, this->G2_factor.data(), n); - - if(this->timing) { - t_stop = steady_clock::now(); - chol2_dur = duration_cast(t_stop - t_start).count(); - t_start = steady_clock::now(); - } - - // R = R2 * R1 - lapack::lacpy(MatrixType::Upper, n, n, R, ldr, R_temp, n); - blas::trmm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, n, n, (T)1.0, G, n, R_temp, n); - lapack::lacpy(MatrixType::Upper, n, n, R_temp, n, R, ldr); - - // M = M * R2^{-1} - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, n, n, (T)1.0, G, n, M, n); - - if(this->timing) { - t_stop = steady_clock::now(); - upd2_dur = duration_cast(t_stop - t_start).count(); - } - - //================================================================ - // Iteration 3: Cholesky QR - //================================================================ - // Blocked Gram: G3 = M^T * A^T * A * M where M = R1^{-1} * R2^{-1} - long fwd3_accum = 0, adj3_accum = 0, gemm3_accum = 0; - for (int64_t j = 0; j < n; j += b_eff) { - int64_t b_j = std::min(b_eff, n - j); - - // W = A * M[:, j:j+b] - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, b_j, n, (T)1.0, M + j * n, n, (T)0.0, A_temp, m); - if(this->timing) { t_stop = steady_clock::now(); fwd3_accum += duration_cast(t_stop - t_start).count(); } - - // Z = A^T * W - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, - n, b_j, m, (T)1.0, A_temp, m, (T)0.0, Z_buf, n); - if(this->timing) { t_stop = steady_clock::now(); adj3_accum += duration_cast(t_stop - t_start).count(); } - - // G[:, j:j+b] = M^T * Z - if(this->timing) t_start = steady_clock::now(); - blas::gemm(Layout::ColMajor, Op::Trans, Op::NoTrans, - n, b_j, n, (T)1.0, M, n, Z_buf, n, (T)0.0, G + j * n, n); - if(this->timing) { t_stop = steady_clock::now(); gemm3_accum += duration_cast(t_stop - t_start).count(); } - } - - if(this->timing) { - fwd3_dur = fwd3_accum; - adj3_dur = adj3_accum; - gemm3_dur = gemm3_accum; - t_start = steady_clock::now(); - } - - // Zero lower triangle, Cholesky: G = R3^T * R3 - if (n > 1) - lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &G[1], n); - if (lapack::potrf(Uplo::Upper, n, G, n)) { - delete[] G; delete[] R_temp; delete[] M; - delete[] A_temp; delete[] Z_buf; - return 3; - } - - // Save G3 factor - this->G3_factor.resize(n * n, (T)0.0); - lapack::lacpy(MatrixType::Upper, n, n, G, n, this->G3_factor.data(), n); - - if(this->timing) { - t_stop = steady_clock::now(); - chol3_dur = duration_cast(t_stop - t_start).count(); - t_start = steady_clock::now(); - } - - // R = R3 * R - lapack::lacpy(MatrixType::Upper, n, n, R, ldr, R_temp, n); - blas::trmm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, n, n, (T)1.0, G, n, R_temp, n); - lapack::lacpy(MatrixType::Upper, n, n, R_temp, n, R, ldr); - - if(this->timing) { - t_stop = steady_clock::now(); - upd3_dur = duration_cast(t_stop - t_start).count(); - } - - //================================================================ - // Test mode: materialize Q = A * R^{-1} = A * M * R3^{-1} - //================================================================ - if(this->test_mode) { - if(this->timing) - t_start = steady_clock::now(); - - // M currently holds R1^{-1} R2^{-1}; update to R^{-1} = R1^{-1} R2^{-1} R3^{-1} - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, n, n, (T)1.0, G, n, M, n); - - // Materialize Q = A * M in blocks - T* Q_buf = new T[m * n](); - for (int64_t j = 0; j < n; j += b_eff) { - int64_t b_j = std::min(b_eff, n - j); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, b_j, n, (T)1.0, M + j * n, n, (T)0.0, Q_buf + j * m, m); - } - + // sCholQR3 = three cholqr_iterate passes; iter 1 carries the shift, + // iters 2-3 use shift_factor_iter23. + // + // The first-pass shift defaults to the paper's prescription + // s = 11*eps*n*trace(G) (FukayaEtAl2020, c = 11), resolved here because + // it needs n. RANDLAPACK_SCHOLQR3_SHIFT=eps selects the legacy smaller + // s = eps*trace(G) instead; a caller-set shift_factor_iter1 >= 0 wins + // over both. + long it[15] = {0}; + const T eps_T = std::numeric_limits::epsilon(); + const T sf1 = (this->shift_factor_iter1 >= T(0)) + ? this->shift_factor_iter1 + : (scholqr3_eps_shift() ? eps_T : T(11) * T(n) * eps_T); + int info = cholqr_iterate( + A, R, ldr, this->block_size, /*num_iters=*/3, + sf1, this->shift_factor_iter23, + this->max_retries, this->shift_growth, this->timing, + this->timing ? it : nullptr, &this->n_chol_retries, + this->chol_applied_shifts, this->chol_gram_traces); + // 1/2/3 = the 1-based pass that failed (retry exhaustion, singular + // preconditioner, non-finite shift, or invalid input; see stderr). + if (info != 0) return info; + + // Test mode: materialize Q = A * R^{-1} (outside timing region). + if (this->test_mode) { + if (this->timing) t0 = steady_clock::now(); + + T* Q_buf = new T[m * n]; + RandLAPACK::materialize_Q_from_R(A, R, ldr, m, n, b_eff, Q_buf, m); this->Q_rows = m; this->Q_cols = n; + // The class owns Q (the destructor frees it), so release any buffer from a + // previous call() before taking ownership of this one. + delete[] this->Q; this->Q = Q_buf; - if(this->timing) { - t_stop = steady_clock::now(); - q_mat_dur = duration_cast(t_stop - t_start).count(); - } + if (this->timing) { t1 = steady_clock::now(); q_mat_dur = duration_cast(t1 - t0).count(); } } - //================================================================ + // ============================================================ // Finalize timing - //================================================================ - if(this->timing) { + // ============================================================ + if (this->timing) { total_t_stop = steady_clock::now(); - total_dur = duration_cast(total_t_stop - total_t_start).count(); - - // Subtract Q materialization from total (test overhead, not algorithmic cost) - total_dur -= q_mat_dur; - - long rest_dur = total_dur - (alloc_dur + - fwd1_dur + adj1_dur + chol1_dur + upd1_dur + - fwd2_dur + adj2_dur + gemm2_dur + chol2_dur + upd2_dur + - fwd3_dur + adj3_dur + gemm3_dur + chol3_dur + upd3_dur); - - // 18 entries - this->times = {alloc_dur, - fwd1_dur, adj1_dur, chol1_dur, upd1_dur, - fwd2_dur, adj2_dur, gemm2_dur, chol2_dur, upd2_dur, - fwd3_dur, adj3_dur, gemm3_dur, chol3_dur, upd3_dur, + long total_dur = duration_cast(total_t_stop - total_t_start).count() - q_mat_dur; + long iters_sum = 0; + for (int i = 0; i < 15; ++i) iters_sum += it[i]; + long rest_dur = total_dur - iters_sum; + // it groups [fwd,adj,gemm,chol,upd] per pass; iter-1 gemm/upd are 0. + this->times = {0L, + it[0], it[1], it[3], it[4], // fwd1, adj1, chol1, upd1 + it[5], it[6], it[7], it[8], it[9], // fwd2, adj2, gemm2, chol2, upd2 + it[10], it[11], it[12], it[13], it[14], // fwd3, adj3, gemm3, chol3, upd3 q_mat_dur, rest_dur, total_dur}; } - // Cleanup - delete[] G; - delete[] R_temp; - delete[] M; - delete[] A_temp; - delete[] Z_buf; - return 0; } }; -/// Non-blocked (basic) sCholQR3 algorithm for computing QR factorization via linear operators. -/// -/// Matches the standard sCholQR3 pseudocode from Fukaya et al. (SISC, 2020) exactly: -/// 1. Compute G1 = A^T A via linop, add shift, Cholesky → R1 -/// 2. Materialize Q = A * R1^{-1} via linop -/// 3. Iterations 2-3: G = Q^T Q via dense syrk, Cholesky, Q *= R_k^{-1} via dense trsm -/// -/// Accesses the linear operator exactly 3 times: -/// - NoTrans: W = A * I (materialization for Gram computation) -/// - Trans: G1 = A^T * W (Gram matrix) -/// - NoTrans: Q = A * R1^{-1} (first Q-factor) -/// -/// After the first Q-factor, iterations 2-3 use dense syrk on Q (no further linop calls). -/// This is theoretically distinct from sCholQR3_linops (fully-blocked), which recomputes -/// each Gram through the linop and never materializes the m x n operator product. -/// -/// Peak memory: O(m*n + n^2) — Q is explicitly stored as m x n dense. -/// -/// Reference: Shifted Cholesky QR from Fukaya et al. (SISC, 2020). +/// Non-blocked (basic) sCholQR3: algorithmically identical to sCholQR3_linops with +/// block_size = 0: all three iterations route through cholqr_primitive on the linop +/// (no materialized-Q / dense-syrk shortcut). It exists only as a separate analytic- +/// memory accounting case; the heavy work is the same per-iteration linop Gram. /// template class sCholQR3_linops_basic { @@ -467,35 +207,38 @@ class sCholQR3_linops_basic { bool timing; bool test_mode; - T eps; - // Q-factor for test mode (only allocated if test_mode = true) T* Q; int64_t Q_rows; int64_t Q_cols; - // Individual Cholesky factors from each iteration (n x n upper triangular). - std::vector G1_factor; - std::vector G2_factor; - std::vector G3_factor; - - // Timing breakdown (15 entries): - // [0] alloc - buffer allocation - // [1] fwd1 - NoTrans: W = A * I (m x n) - // [2] adj1 - Trans: G = A^T * W (n x n) - // [3] chol1 - Iter 1 Cholesky - // [4] trsm1 - M = R1^{-1} (n x n trsm) - // [5] fwd_q - NoTrans: Q = A * M (m x n) - // [6] syrk2 - G = Q^T Q - // [7] chol2 - Iter 2 Cholesky - // [8] upd2 - Q *= R2^{-1}, R = R2*R1 - // [9] syrk3 - G = Q^T Q - // [10] chol3 - Iter 3 Cholesky - // [11] upd3 - R = R3*R - // [12] q_mat - test mode: Q_buf *= R3^{-1} (m x n trsm), 0 otherwise - // [13] rest - unaccounted time - // [14] total - wall-clock total + // Timing layout (15 entries, kept for matlab CSV-column compatibility): + // [0] alloc [1] fwd1 [2] adj1 [3] chol1 [4] trsm1=0 [5] fwd_q=0 + // [6] syrk2 [7] chol2 [8] upd2 + // [9] syrk3 [10] chol3 [11] upd3 + // [12] q_mat [13] rest [14] total + // (Post-refactor slots 4, 5, 6, 9 stay 0 because the primitives don't expose + // syrk vs adj/fwd as separate signals; the heavy lifters are folded into + // fwd/adj from blocked_preconditioned_gram and into chol from potrf.) std::vector times; + /// Total measured wall-clock (microseconds) of the last call(), or -1 if timing + /// was off. Every driver in this family packs the total as the LAST times[] entry, + /// but the entry COUNT differs per driver (6 / 11 / 15 / 18). Callers used to hard- + /// code that index (times[5], times[10], times[14], times[17]), so adding or + /// removing one slot silently wrote the wrong number into every CSV with no compile + /// error. Read the total through here instead. + long total_us() const { return times.empty() ? -1L : times.back(); } + + // Adaptive shift policy, shared with sCholQR3_linops. + T shift_factor_iter1; + T shift_factor_iter23; + int max_retries; + T shift_growth; + int n_chol_retries = 0; ///< shift retries used on the last call (0 = clean) + /// Per-pass shift record from the last call (passes 1-3): absolute shift the + /// successful potrf carried (0 = unshifted) and that pass's Gram trace. + T chol_applied_shifts[3] = {T(0), T(0), T(0)}; + T chol_gram_traces[3] = {T(0), T(0), T(0)}; sCholQR3_linops_basic( bool time_subroutines, @@ -503,11 +246,15 @@ class sCholQR3_linops_basic { bool enable_test_mode = false ) { timing = time_subroutines; - eps = ep; + (void)ep; // kept in the signature for call-site compatibility; unused test_mode = enable_test_mode; Q = nullptr; Q_rows = 0; Q_cols = 0; + shift_factor_iter1 = T(-1); // < 0: resolve default (11*n*eps, or eps via env) at call time + shift_factor_iter23 = T(0); // iters 2-3 unshifted (Fukaya); retry covers genuine non-PD + max_retries = -1; // unbounded retries (no ceiling), consistent with CholQR/CholQR2 + shift_growth = T(10); } ~sCholQR3_linops_basic() { @@ -516,269 +263,73 @@ class sCholQR3_linops_basic { } } - /// Computes the QR factorization A = QR using shifted Cholesky QR3 (basic variant). - /// - /// @param[in] A - /// The m-by-n linear operator (m and n read from A.n_rows, A.n_cols). - /// - /// @param[out] R - /// Pre-allocated n-by-n buffer. On exit, stores the upper-triangular - /// R factor. Zero entries are not compressed. - /// - /// @param[in] ldr - /// Leading dimension of R. - /// - /// @return = 0: successful exit + // Non-blocked sCholQR3 expressed via the shared primitives. + // Algorithmically identical to sCholQR3_linops with block_size=0; the + // distinction is now purely the analytic-memory accounting (no R_pre / + // P_prev / Z_buf reuse across iters since the primitives allocate their + // own G internally per call). Diagnostic prints from cholqr_primitive + // surface here too. template int call( GLO& A, T* R, int64_t ldr ) { - ///--------------------TIMING VARS--------------------/ - steady_clock::time_point t_start, t_stop; - steady_clock::time_point total_t_start, total_t_stop; - long alloc_dur = 0; - long fwd1_dur = 0, adj1_dur = 0, chol1_dur = 0, trsm1_dur = 0, fwd_q_dur = 0; - long syrk2_dur = 0, chol2_dur = 0, upd2_dur = 0; - long syrk3_dur = 0, chol3_dur = 0, upd3_dur = 0; - long q_mat_dur = 0, total_dur = 0; - - if(this->timing) - total_t_start = steady_clock::now(); + steady_clock::time_point t0, t1, total_t_start, total_t_stop; + long q_mat_dur = 0; + + if (this->timing) total_t_start = steady_clock::now(); int64_t m = A.n_rows; int64_t n = A.n_cols; - - if(this->timing) - t_start = steady_clock::now(); - - // ---- Allocate buffers ---- - // Q_buf: m x n — materialized operator, updated in-place through iterations - T* Q_buf = new T[m * n]; - - // G: n x n Gram matrix / Cholesky workspace (zero-init for lower triangle) - T* G = new T[n * n](); - - // R_temp: n x n workspace for R accumulation via trmm - T* R_temp = new T[n * n](); - - // M: n x n — starts as identity, becomes R1^{-1} for Q materialization - T* M = new T[n * n](); - RandLAPACK::util::eye(n, n, M); - - if(this->timing) { - t_stop = steady_clock::now(); - alloc_dur = duration_cast(t_stop - t_start).count(); - } - - //================================================================ - // Iteration 1: Shifted Cholesky QR - //================================================================ - // Gram: G1 = A^T * A via linop (2 of 3 total linop accesses) - - // Linop access 1: W = A * I (NoTrans, materializes operator as m x n dense) - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, n, n, (T)1.0, M, n, (T)0.0, Q_buf, m); - if(this->timing) { t_stop = steady_clock::now(); fwd1_dur = duration_cast(t_stop - t_start).count(); } - - // Linop access 2: G1 = A^T * W (Trans, n x n Gram matrix) - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, - n, n, m, (T)1.0, Q_buf, m, (T)0.0, G, n); - if(this->timing) { t_stop = steady_clock::now(); adj1_dur = duration_cast(t_stop - t_start).count(); } - - // Compute shift from ||A||_F^2 = trace(G) - T norm_A_sq = 0; - for (int64_t i = 0; i < n; ++i) - norm_A_sq += G[i * (n + 1)]; - T shift = 11 * std::numeric_limits::epsilon() * n * norm_A_sq; - - // Add shift to diagonal: G = G + shift * I - for (int64_t i = 0; i < n; ++i) - G[i * (n + 1)] += shift; - - if(this->timing) { - t_start = steady_clock::now(); - } - - // Zero lower triangle, Cholesky: G = R1^T * R1 - if (n > 1) - lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &G[1], n); - if (lapack::potrf(Uplo::Upper, n, G, n)) { - delete[] Q_buf; delete[] G; delete[] R_temp; delete[] M; - return 1; - } - - // Save G1 factor - this->G1_factor.resize(n * n, (T)0.0); - lapack::lacpy(MatrixType::Upper, n, n, G, n, this->G1_factor.data(), n); - - // Initialize R = R1 - lapack::lacpy(MatrixType::Upper, n, n, G, n, R, ldr); - - if(this->timing) { - t_stop = steady_clock::now(); - chol1_dur = duration_cast(t_stop - t_start).count(); - } - - // Compute M = I * R1^{-1} = R1^{-1} - if(this->timing) t_start = steady_clock::now(); - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, n, n, (T)1.0, G, n, M, n); - if(this->timing) { t_stop = steady_clock::now(); trsm1_dur = duration_cast(t_stop - t_start).count(); } - - // Linop access 3: Q_buf = A * M = A * R1^{-1} (NoTrans, m x n) - if(this->timing) t_start = steady_clock::now(); - A(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, - m, n, n, (T)1.0, M, n, (T)0.0, Q_buf, m); - if(this->timing) { t_stop = steady_clock::now(); fwd_q_dur = duration_cast(t_stop - t_start).count(); } - - //================================================================ - // Iteration 2: Cholesky QR (dense syrk on Q_buf) - //================================================================ - if(this->timing) - t_start = steady_clock::now(); - - // G2 = Q_buf^T * Q_buf (dense syrk, upper triangle only) - blas::syrk(Layout::ColMajor, Uplo::Upper, Op::Trans, - n, m, (T)1.0, Q_buf, m, (T)0.0, G, n); - - if(this->timing) { - t_stop = steady_clock::now(); - syrk2_dur = duration_cast(t_stop - t_start).count(); - t_start = steady_clock::now(); - } - - // Zero lower triangle, Cholesky: G = R2^T * R2 - if (n > 1) - lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &G[1], n); - if (lapack::potrf(Uplo::Upper, n, G, n)) { - delete[] Q_buf; delete[] G; delete[] R_temp; delete[] M; - return 2; - } - - // Save G2 factor - this->G2_factor.resize(n * n, (T)0.0); - lapack::lacpy(MatrixType::Upper, n, n, G, n, this->G2_factor.data(), n); - - if(this->timing) { - t_stop = steady_clock::now(); - chol2_dur = duration_cast(t_stop - t_start).count(); - t_start = steady_clock::now(); - } - - // Q_buf *= R2^{-1} (m x n trsm — update Q in-place) - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, m, n, (T)1.0, G, n, Q_buf, m); - - // R = R2 * R1 - lapack::lacpy(MatrixType::Upper, n, n, R, ldr, R_temp, n); - blas::trmm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, n, n, (T)1.0, G, n, R_temp, n); - lapack::lacpy(MatrixType::Upper, n, n, R_temp, n, R, ldr); - - if(this->timing) { - t_stop = steady_clock::now(); - upd2_dur = duration_cast(t_stop - t_start).count(); - } - - //================================================================ - // Iteration 3: Cholesky QR (dense syrk on Q_buf) - //================================================================ - if(this->timing) - t_start = steady_clock::now(); - - // G3 = Q_buf^T * Q_buf (dense syrk, upper triangle only) - blas::syrk(Layout::ColMajor, Uplo::Upper, Op::Trans, - n, m, (T)1.0, Q_buf, m, (T)0.0, G, n); - - if(this->timing) { - t_stop = steady_clock::now(); - syrk3_dur = duration_cast(t_stop - t_start).count(); - t_start = steady_clock::now(); - } - - // Zero lower triangle, Cholesky: G = R3^T * R3 - if (n > 1) - lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &G[1], n); - if (lapack::potrf(Uplo::Upper, n, G, n)) { - delete[] Q_buf; delete[] G; delete[] R_temp; delete[] M; - return 3; - } - - // Save G3 factor - this->G3_factor.resize(n * n, (T)0.0); - lapack::lacpy(MatrixType::Upper, n, n, G, n, this->G3_factor.data(), n); - - if(this->timing) { - t_stop = steady_clock::now(); - chol3_dur = duration_cast(t_stop - t_start).count(); - t_start = steady_clock::now(); - } - - // R = R3 * R - lapack::lacpy(MatrixType::Upper, n, n, R, ldr, R_temp, n); - blas::trmm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, n, n, (T)1.0, G, n, R_temp, n); - lapack::lacpy(MatrixType::Upper, n, n, R_temp, n, R, ldr); - - if(this->timing) { - t_stop = steady_clock::now(); - upd3_dur = duration_cast(t_stop - t_start).count(); - } - - //================================================================ - // Test mode: Q = Q_buf * R3^{-1} - //================================================================ - if(this->test_mode) { - if(this->timing) - t_start = steady_clock::now(); - - // Q_buf currently holds A * R1^{-1} * R2^{-1} - // Apply R3^{-1}: Q_buf = Q_buf * R3^{-1} = A * R^{-1} = Q - blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, - Diag::NonUnit, m, n, (T)1.0, G, n, Q_buf, m); - + int64_t b_eff = n; // non-blocked (block_size = 0 → b_eff = n) + + // Non-blocked sCholQR3 = three cholqr_iterate passes with block_size = 0. + long it[15] = {0}; + // Same shift-default resolution as the blocked variant above. + const T eps_T = std::numeric_limits::epsilon(); + const T sf1 = (this->shift_factor_iter1 >= T(0)) + ? this->shift_factor_iter1 + : (scholqr3_eps_shift() ? eps_T : T(11) * T(n) * eps_T); + int info = cholqr_iterate( + A, R, ldr, /*block_size=*/0, /*num_iters=*/3, + sf1, this->shift_factor_iter23, + this->max_retries, this->shift_growth, this->timing, + this->timing ? it : nullptr, &this->n_chol_retries, + this->chol_applied_shifts, this->chol_gram_traces); + // 1/2/3 = the 1-based pass that failed (retry exhaustion, singular + // preconditioner, non-finite shift, or invalid input; see stderr). + if (info != 0) return info; + + // ---- Test mode: materialize Q = A * R^{-1} via blocked linop call ---- + if (this->test_mode) { + if (this->timing) t0 = steady_clock::now(); + T* Q_buf = new T[m * n]; + RandLAPACK::materialize_Q_from_R(A, R, ldr, m, n, b_eff, Q_buf, m); this->Q_rows = m; this->Q_cols = n; - this->Q = Q_buf; // Take ownership of Q_buf - - if(this->timing) { - t_stop = steady_clock::now(); - q_mat_dur = duration_cast(t_stop - t_start).count(); - } + // The class owns Q (the destructor frees it), so release any buffer from a + // previous call() before taking ownership of this one. + delete[] this->Q; + this->Q = Q_buf; + if (this->timing) { t1 = steady_clock::now(); q_mat_dur = duration_cast(t1 - t0).count(); } } - //================================================================ - // Finalize timing - //================================================================ - if(this->timing) { + if (this->timing) { total_t_stop = steady_clock::now(); - total_dur = duration_cast(total_t_stop - total_t_start).count(); - - // Subtract Q materialization from total (test overhead, not algorithmic cost) - total_dur -= q_mat_dur; - - long rest_dur = total_dur - (alloc_dur + fwd1_dur + adj1_dur + chol1_dur + trsm1_dur + fwd_q_dur + - syrk2_dur + chol2_dur + upd2_dur + - syrk3_dur + chol3_dur + upd3_dur); - - // 15 entries - this->times = {alloc_dur, fwd1_dur, adj1_dur, chol1_dur, trsm1_dur, fwd_q_dur, - syrk2_dur, chol2_dur, upd2_dur, - syrk3_dur, chol3_dur, upd3_dur, + long total_dur = duration_cast(total_t_stop - total_t_start).count() - q_mat_dur; + long iters_sum = 0; + for (int i = 0; i < 15; ++i) iters_sum += it[i]; + long rest_dur = total_dur - iters_sum; + // Basic layout folds each iter's fwd+adj+gemm into its chol slot. + long chol2 = it[8] + it[5] + it[6] + it[7]; + long chol3 = it[13] + it[10] + it[11] + it[12]; + this->times = {0L, it[0], it[1], it[3], + /*trsm1=*/0L, /*fwd_q=*/0L, + /*syrk2=*/0L, chol2, it[9], + /*syrk3=*/0L, chol3, it[14], q_mat_dur, rest_dur, total_dur}; } - - // Cleanup - delete[] G; - delete[] R_temp; - delete[] M; - if(!this->test_mode) - delete[] Q_buf; - return 0; } }; diff --git a/RandLAPACK/linops/rl_linops.hh b/RandLAPACK/linops/rl_linops.hh index b981f2c35..0a66f2835 100644 --- a/RandLAPACK/linops/rl_linops.hh +++ b/RandLAPACK/linops/rl_linops.hh @@ -14,5 +14,9 @@ #include "rl_dense_linop.hh" #include "rl_sparse_linop.hh" #include "rl_composite_linop.hh" +#include "rl_power_linop.hh" +#include "rl_transposed_linop.hh" +#include "rl_scaled_identity_linop.hh" +#include "rl_vstack_linop.hh" #include "rl_sym_linops.hh" #include "rl_materialize.hh" diff --git a/RandLAPACK/linops/rl_power_linop.hh b/RandLAPACK/linops/rl_power_linop.hh new file mode 100644 index 000000000..46685a364 --- /dev/null +++ b/RandLAPACK/linops/rl_power_linop.hh @@ -0,0 +1,179 @@ +#pragma once + +// Public API: PowerOp: implicit j-th power of a square linear operator. + +#include "rl_concepts.hh" +#include "rl_blaspp.hh" +#include "rl_exceptions.hh" + +#include +#include +#include +#include + + +namespace RandLAPACK::linops { + +/*********************************************************/ +/* */ +/* PowerOp */ +/* */ +/*********************************************************/ +// Generic LinearOperator wrapper that represents A^j for a square base operator A. +// +// Template parameter: +// InnerOp - Square base operator satisfying LinearOperator concept (dense, sparse, +// composite, sparse-solver-inverse, etc.) +// +// Strategy: +// Each application chains j calls to the base operator with two ping-pong scratch +// buffers. A^j is never materialized. +// +// j == 1: single base call, no scratch. +// j >= 2: one scratch for the first apply; a second scratch for the middle applies +// (j == 2 only allocates the first). The final apply writes to C and +// respects the user's alpha/beta. +// +// Restrictions: +// - base.n_rows must equal base.n_cols. PowerOp is only well-defined for square base. +// - Side::Left only (square A^j on the left of B). Side::Right could be added by +// mirroring the loop, but no current consumer needs it. +// +// Op::Trans semantics: +// trans_A == Op::Trans applies (A^T)^j == (A^j)^T: each iteration dispatches the +// base op with Op::Trans. Intermediate scratch dispatches always use Op::NoTrans. +// +template +struct PowerOp { + using T = typename InnerOp::scalar_t; + using scalar_t = T; + + InnerOp& base; + const int j; + const int64_t n_rows; + const int64_t n_cols; + + PowerOp(InnerOp& base_op, int power) + : base(base_op), j(power), + n_rows(base_op.n_rows), n_cols(base_op.n_cols) + { + randlapack_require(base.n_rows == base.n_cols) + << "PowerOp: base must be square, got n_rows=" << base.n_rows + << " n_cols=" << base.n_cols; + randlapack_require(power >= 1) + << "PowerOp: power=" << power << " must be >= 1 (identity j=0 not supported; caller can lacpy)"; + } + + // Concept-required 12-arg overload (no Side); delegates to Side::Left. + void operator()( + Layout layout, Op trans_A, Op trans_B, + int64_t m, int64_t n, int64_t k, + T alpha, const T* B, int64_t ldb, + T beta, T* C, int64_t ldc) + { + (*this)(Side::Left, layout, trans_A, trans_B, + m, n, k, alpha, B, ldb, beta, C, ldc); + } + + // C := alpha * (base^j)^{trans_A} * op_{trans_B}(B) + beta * C + // + // Since base is square (N x N), m == k == N for any valid Side::Left call. + // op_{trans_B}(B) has shape N x n; C has shape N x n. + void operator()( + Side side, Layout layout, + Op trans_A, Op trans_B, + int64_t m, int64_t n, int64_t k, + T alpha, const T* B, int64_t ldb, + T beta, T* C, int64_t ldc) + { + randlapack_require(side == Side::Left) << "PowerOp supports Side::Left only"; + randlapack_require(m == n_rows) << "PowerOp: m=" << m << " must equal n_rows=" << n_rows; + randlapack_require(k == n_rows) << "PowerOp: k=" << k << " must equal n_rows=" << n_rows; + + // side is fixed to Left above, so base is dispatched through the concept- + // required 12-arg overload (no Side param needed); this is the form every + // LinearOperator is guaranteed to provide. + if (j == 1) { + base(layout, trans_A, trans_B, + m, n, k, alpha, B, ldb, beta, C, ldc); + return; + } + + // j >= 2: ping-pong scratch. + // Layout-aware leading dimension: m for ColMajor (m x n stored), n for RowMajor. + int64_t ldt = (layout == Layout::ColMajor) ? m : n; + T* buf_a = new T[(size_t)m * (size_t)n](); + T* buf_b = (j >= 3) ? new T[(size_t)m * (size_t)n]() : nullptr; + + // 1st apply: buf_a := base^{trans_A} * op_{trans_B}(B) + base(layout, trans_A, trans_B, + m, n, k, (T)1.0, B, ldb, (T)0.0, buf_a, ldt); + + // Middle applies (only run when j >= 3): ping-pong buf_a <-> buf_b. + T* in_buf = buf_a; + T* out_buf = buf_b; + for (int it = 1; it < j - 1; ++it) { + base(layout, trans_A, Op::NoTrans, + m, n, m, (T)1.0, in_buf, ldt, (T)0.0, out_buf, ldt); + std::swap(in_buf, out_buf); + } + + // Last apply writes to C with the user's alpha/beta. + base(layout, trans_A, Op::NoTrans, + m, n, m, alpha, in_buf, ldt, beta, C, ldc); + + delete[] buf_a; + if (buf_b) delete[] buf_b; + } + + // SkOp overload: materialize S as a dense matrix, then delegate to the dense apply. + // Square base means op_{trans_A}(base^j) is square N x N, so op(S) must be N x n, + // i.e. S is k x n or n x k. Side::Left only (checked before any allocation). + template + void operator()( + Side side, Layout layout, + Op trans_A, Op trans_S, + int64_t m, int64_t n, int64_t k, + T alpha, SkOp& S, + T beta, T* C, int64_t ldc) + { + randlapack_require(side == Side::Left) << "PowerOp SkOp overload supports Side::Left only"; + + // SparseSkOp materialization only supports ColMajor (see below). Checked + // before S_dense is allocated so the reject path cannot leak it. + if constexpr (!std::is_same_v) { + randlapack_require(layout == Layout::ColMajor) + << "PowerOp SkOp overload materializes SparseSkOp in ColMajor only, got RowMajor"; + } + + int64_t S_rows = S.n_rows; + int64_t S_cols = S.n_cols; + int64_t lds = (layout == Layout::ColMajor) ? S_rows : S_cols; + + T* S_dense = new T[(size_t)S_rows * (size_t)S_cols]; + + if constexpr (std::is_same_v) { + // Materialize directly in the caller's layout: no sketch-by-identity + // GEMM, no S_cols^2 identity buffer. + RandBLAS::fill_dense_unpacked(layout, S.dist, S_rows, S_cols, 0, 0, S_dense, S.seed_state); + } else { + T* I_block = new T[(size_t)S_cols * (size_t)S_cols](); + for (int64_t i = 0; i < S_cols; ++i) I_block[i + i * S_cols] = (T)1.0; + RandBLAS::sketch_general( + Layout::ColMajor, Op::NoTrans, Op::NoTrans, + S_rows, S_cols, S_cols, + (T)1.0, S, I_block, S_cols, + (T)0.0, S_dense, S_rows); + delete[] I_block; + } + + // Delegate to the dense overload's concept-required 12-arg form (side is + // fixed to Left above). + (*this)(layout, trans_A, trans_S, + m, n, k, alpha, S_dense, lds, beta, C, ldc); + + delete[] S_dense; + } +}; + +} // namespace RandLAPACK::linops diff --git a/RandLAPACK/linops/rl_scaled_identity_linop.hh b/RandLAPACK/linops/rl_scaled_identity_linop.hh new file mode 100644 index 000000000..2aafcc9d9 --- /dev/null +++ b/RandLAPACK/linops/rl_scaled_identity_linop.hh @@ -0,0 +1,93 @@ +#pragma once + +// Public API: ScaledIdentityOp: matrix-free scaled identity mu*I_n. + +#include "rl_exceptions.hh" +#include "rl_blaspp.hh" + +#include + + +namespace RandLAPACK::linops { + +/*********************************************************/ +/* */ +/* ScaledIdentityOp */ +/* */ +/*********************************************************/ +// Matrix-free n x n scaled identity operator mu * I_n. +// +// Its main use is as the bottom block of a VStackOp to build the regularized +// augmented operator A_hat = [A; mu*I], whose Gram is A^T A + mu^2 I. A +// Cholesky-QR of A_hat therefore yields R = chol(A^T A + mu^2 I), a regularized +// right preconditioner that is well defined even when A is rank-deficient or +// extremely ill-conditioned (see rl_iter_refine_lsq.hh). +// +// mu*I is symmetric, so Op::NoTrans and Op::Trans behave identically. Only +// Side::Left and trans_B == Op::NoTrans are supported: that is all the +// Cholesky-QR Gram path, IterRefineLSQ, and the orthogonality check ever use. +template +struct ScaledIdentityOp { + using scalar_t = T; + const int64_t n_rows; // = n + const int64_t n_cols; // = n + const T mu; + + ScaledIdentityOp(int64_t n, T mu_) + : n_rows(n), n_cols(n), mu(mu_) {} + + // Concept-required 12-arg overload (no Side); delegates to Side::Left. + void operator()( + Layout layout, Op trans_self, Op trans_B, + int64_t m, int64_t n, int64_t k, + T alpha, const T* B, int64_t ldb, + T beta, T* C, int64_t ldc) + { + (*this)(Side::Left, layout, trans_self, trans_B, + m, n, k, alpha, B, ldb, beta, C, ldc); + } + + // C := alpha * (mu I) * B + beta * C (mu I is symmetric, so trans_self is moot). + // The identity action requires the contracted dim k to equal the output row + // dim m (square identity), so C[i,j] = alpha*mu*B[i,j] + beta*C[i,j]. + void operator()( + Side side, Layout layout, + Op trans_self, Op trans_B, + int64_t m, int64_t n, int64_t k, + T alpha, const T* B, int64_t ldb, + T beta, T* C, int64_t ldc) + { + (void)trans_self; // mu*I is symmetric. + randlapack_require(side == Side::Left) << "ScaledIdentityOp supports Side::Left only"; + randlapack_require(trans_B == Op::NoTrans) << "ScaledIdentityOp supports trans_B == NoTrans only"; + randlapack_require(k == m) << "ScaledIdentityOp: contracted dim k=" << k + << " must equal output rows m=" << m << " (square identity)"; + randlapack_require(m == n_rows) << "ScaledIdentityOp: m=" << m + << " must equal n_rows=" << n_rows; + + const T am = alpha * mu; + if (layout == Layout::ColMajor) { + for (int64_t j = 0; j < n; ++j) { + const T* bcol = B + j * ldb; + T* ccol = C + j * ldc; + if (beta == (T)0) { + for (int64_t i = 0; i < m; ++i) ccol[i] = am * bcol[i]; + } else { + for (int64_t i = 0; i < m; ++i) ccol[i] = beta * ccol[i] + am * bcol[i]; + } + } + } else { // RowMajor + for (int64_t i = 0; i < m; ++i) { + const T* brow = B + i * ldb; + T* crow = C + i * ldc; + if (beta == (T)0) { + for (int64_t j = 0; j < n; ++j) crow[j] = am * brow[j]; + } else { + for (int64_t j = 0; j < n; ++j) crow[j] = beta * crow[j] + am * brow[j]; + } + } + } + } +}; + +} // namespace RandLAPACK::linops diff --git a/RandLAPACK/linops/rl_transposed_linop.hh b/RandLAPACK/linops/rl_transposed_linop.hh new file mode 100644 index 000000000..5bd6fc756 --- /dev/null +++ b/RandLAPACK/linops/rl_transposed_linop.hh @@ -0,0 +1,111 @@ +#pragma once + +// Public API: TransposedOp: implicit transpose view of a LinearOperator. + +#include "rl_concepts.hh" +#include "rl_blaspp.hh" + +#include +#include + + +namespace RandLAPACK::linops { + +/*********************************************************/ +/* */ +/* TransposedOp */ +/* */ +/*********************************************************/ +// Generic LinearOperator wrapper that represents A^T for an inner operator A. +// +// Template parameter: +// InnerOp - Any type satisfying the LinearOperator concept (dense, sparse, +// composite, solver-based, PowerOp, another TransposedOp, ...). +// InnerOp must implement operator() with both Op::NoTrans and +// Op::Trans dispatch on its first matrix argument (this is part +// of the LinearOperator concept). +// +// Semantics: +// TransposedOp simply flips the user-supplied trans_A flag before delegating +// to the inner op. No data is materialized; this is a zero-cost view. +// +// user calls T_op(..., trans_A = NoTrans, ...) → base(..., Trans, ...) +// user calls T_op(..., trans_A = Trans, ...) → base(..., NoTrans, ...) +// +// n_rows and n_cols are swapped from base so non-square wrapping works +// (CompositeOperator etc. read these for dimension checks). +// +// Why this is fully generic: +// The LinearOperator concept already requires both Op::NoTrans and Op::Trans +// dispatch on the first matrix argument. Every implementation that satisfies +// the concept supports being transposed by the right dispatch flag, so this +// wrapper just inverts the mapping. Composite chains, sparse solvers, +// PowerOp, and nested TransposedOps all flow through the same one-line body. +// +template +struct TransposedOp { + using T = typename InnerOp::scalar_t; + using scalar_t = T; + + InnerOp& base; + const int64_t n_rows; // = base.n_cols + const int64_t n_cols; // = base.n_rows + + explicit TransposedOp(InnerOp& base_op) + : base(base_op), n_rows(base_op.n_cols), n_cols(base_op.n_rows) {} + + // Concept-required 12-arg overload (no Side); delegates to Side::Left. + void operator()( + Layout layout, Op trans_self, Op trans_B, + int64_t m, int64_t n, int64_t k, + T alpha, const T* B, int64_t ldb, + T beta, T* C, int64_t ldc) + { + (*this)(Side::Left, layout, trans_self, trans_B, + m, n, k, alpha, B, ldb, beta, C, ldc); + } + + // C := alpha * op_{trans_self}(base^T) * op_{trans_B}(B) + beta * C + // + // op_{NoTrans}(base^T) = base^T, dispatched to base() with Op::Trans. + // op_{Trans}(base^T) = base, dispatched to base() with Op::NoTrans. + // + // side == Side::Left dispatches through the concept-required 12-arg overload + // (every LinearOperator is guaranteed to provide it). side == Side::Right needs + // the extended Side-taking overload, which the LinearOperator concept does not + // guarantee; InnerOp must supply it if this wrapper is ever used on the right. + void operator()( + Side side, Layout layout, + Op trans_self, Op trans_B, + int64_t m, int64_t n, int64_t k, + T alpha, const T* B, int64_t ldb, + T beta, T* C, int64_t ldc) + { + Op base_trans = (trans_self == Op::NoTrans) ? Op::Trans : Op::NoTrans; + if (side == Side::Left) { + base(layout, base_trans, trans_B, + m, n, k, alpha, B, ldb, beta, C, ldc); + } else { + base(side, layout, base_trans, trans_B, + m, n, k, alpha, B, ldb, beta, C, ldc); + } + } + + // SkOp overload: delegate to base by flipping trans_self. SkOp support is not + // part of the LinearOperator concept at all, so InnerOp must provide this + // Side-taking SkOp overload itself; there is no concept-guaranteed fallback. + template + void operator()( + Side side, Layout layout, + Op trans_self, Op trans_S, + int64_t m, int64_t n, int64_t k, + T alpha, SkOp& S, + T beta, T* C, int64_t ldc) + { + Op base_trans = (trans_self == Op::NoTrans) ? Op::Trans : Op::NoTrans; + base(side, layout, base_trans, trans_S, + m, n, k, alpha, S, beta, C, ldc); + } +}; + +} // namespace RandLAPACK::linops diff --git a/RandLAPACK/linops/rl_vstack_linop.hh b/RandLAPACK/linops/rl_vstack_linop.hh new file mode 100644 index 000000000..90ef3d3f3 --- /dev/null +++ b/RandLAPACK/linops/rl_vstack_linop.hh @@ -0,0 +1,164 @@ +#pragma once + +// Public API: VStackOp: vertical concatenation [Top; Bot] of two linear operators. + +#include "rl_exceptions.hh" +#include "rl_concepts.hh" +#include "rl_blaspp.hh" + +#include +#include +#include + + +namespace RandLAPACK::linops { + +/*********************************************************/ +/* */ +/* VStackOp */ +/* */ +/*********************************************************/ +// Vertical (row-wise) concatenation of two linear operators that share a column +// dimension: +// +// A_hat = [ Top ] (Top.n_rows + Bot.n_rows) x n_cols +// [ Bot ] +// +// with Top.n_cols == Bot.n_cols == n_cols. Apply rules: +// +// NoTrans: A_hat * X = [ Top*X ; Bot*X ] (each block written to its rows) +// Trans: A_hat^T * Y = Top^T*Y_1 + Bot^T*Y_2, Y = [Y_1; Y_2] split at Top.n_rows +// +// The headline use is the mu-regularized augmented operator for Cholesky-QR +// preconditioning: A_hat = VStackOp(A, ScaledIdentityOp(mu, n)). Because the +// Gram path forms A_hat^T (A_hat * E) block by block, it yields exactly +// A^T A + mu^2 I with no change to any Cholesky-QR driver, so the resulting R is +// the regularized factor used as a right preconditioner in IterRefineLSQ. +// +// The dense path (Side::Left, trans_B == NoTrans) covers the Cholesky-QR Gram, +// IterRefineLSQ, and the orthogonality check. A sketching overload (Side::Right) +// is also provided so sketch-based drivers (CQRRT) can be handed A_hat directly: +// it is BLOCKED: for each output column block it forms W = A_hat * I_block (an +// (n_rows x b) slice, via this operator's own NoTrans) and sketches that small +// block with the full S, so no (n_rows x d) intermediate is ever materialized and +// S is never partitioned. Works for sparse and dense sketches alike. +template +struct VStackOp { + using T = typename TopOp::scalar_t; + using scalar_t = T; + + TopOp& top; + BotOp& bot; + const int64_t n_rows; // = top.n_rows + bot.n_rows + const int64_t n_cols; // = top.n_cols == bot.n_cols + + /// Block size for the blocked sketch path (0 -> default 256). Caps the width of + /// the (n_rows x b) slice formed per output column block. + int64_t block_size = 0; + + VStackOp(TopOp& top_op, BotOp& bot_op) + : top(top_op), bot(bot_op), + n_rows(top_op.n_rows + bot_op.n_rows), + n_cols(top_op.n_cols) + { + randlapack_require(top_op.n_cols == bot_op.n_cols) + << "VStackOp: top.n_cols=" << top_op.n_cols + << " must match bot.n_cols=" << bot_op.n_cols; + } + + // Concept-required 12-arg overload (no Side); delegates to Side::Left. + void operator()( + Layout layout, Op trans_self, Op trans_B, + int64_t m, int64_t n, int64_t k, + T alpha, const T* B, int64_t ldb, + T beta, T* C, int64_t ldc) + { + (*this)(Side::Left, layout, trans_self, trans_B, + m, n, k, alpha, B, ldb, beta, C, ldc); + } + + void operator()( + Side side, Layout layout, + Op trans_self, Op trans_B, + int64_t m, int64_t n, int64_t k, + T alpha, const T* B, int64_t ldb, + T beta, T* C, int64_t ldc) + { + randlapack_require(side == Side::Left) << "VStackOp supports Side::Left only"; + randlapack_require(trans_B == Op::NoTrans) << "VStackOp supports trans_B == NoTrans only"; + + const int64_t r_top = top.n_rows; + const int64_t r_bot = bot.n_rows; + + // side is fixed to Left above, so top/bot are dispatched through the + // concept-required 12-arg overload (no Side param needed). + if (trans_self == Op::NoTrans) { + // C (n_rows x n) := alpha * [Top; Bot] * B + beta * C. + // Top fills output rows [0, r_top); Bot fills [r_top, r_top + r_bot). + // The two row-blocks are disjoint, so each applies beta to its own region. + randlapack_require(m == n_rows) << "VStackOp NoTrans: m=" << m + << " must equal n_rows=" << n_rows; + randlapack_require(k == n_cols) << "VStackOp NoTrans: k=" << k + << " must equal n_cols=" << n_cols; + const int64_t bot_off = (layout == Layout::ColMajor) ? r_top : r_top * ldc; + top(layout, Op::NoTrans, Op::NoTrans, r_top, n, k, + alpha, B, ldb, beta, C, ldc); + bot(layout, Op::NoTrans, Op::NoTrans, r_bot, n, k, + alpha, B, ldb, beta, C + bot_off, ldc); + } else { + // C (n_cols x n) := alpha * [Top; Bot]^T * B + beta * C + // = alpha * (Top^T * B_top + Bot^T * B_bot) + beta * C, + // where B (n_rows x n) splits row-wise at r_top. + randlapack_require(k == n_rows) << "VStackOp Trans: k=" << k + << " must equal n_rows=" << n_rows; + randlapack_require(m == n_cols) << "VStackOp Trans: m=" << m + << " must equal n_cols=" << n_cols; + const int64_t bot_off = (layout == Layout::ColMajor) ? r_top : r_top * ldb; + // First block sets C (applies beta); second block accumulates (beta = 1). + top(layout, Op::Trans, Op::NoTrans, m, n, r_top, + alpha, B, ldb, beta, C, ldc); + bot(layout, Op::Trans, Op::NoTrans, m, n, r_bot, + alpha, B + bot_off, ldb, (T)1.0, C, ldc); + } + } + + // Blocked sketch: C := alpha * op(S) * [Top; Bot] + beta * C, with S a + // d x n_rows sketching operator (Side::Right means S multiplies on the left of + // this operator). For each output column block we form W = [Top;Bot] * I_block + // (an n_rows x b slice, via this operator's own NoTrans) and sketch it with the + // full S, so the only buffers are O(n_rows x b), no n_rows x d intermediate, + // and S is never partitioned. This is how CQRRT can be handed A_hat directly. + template + void operator()( + Side side, Layout layout, + Op trans_self, Op trans_S, + int64_t m, int64_t n, int64_t k, + T alpha, SkOp& S, T beta, T* C, int64_t ldc) + { + randlapack_require(side == Side::Right) << "VStackOp sketch overload supports Side::Right only"; + randlapack_require(trans_self == Op::NoTrans) << "VStackOp sketch overload supports trans_self == NoTrans only"; + randlapack_require(layout == Layout::ColMajor) << "VStackOp sketch overload supports ColMajor only"; + randlapack_require(k == n_rows) << "VStackOp sketch: k=" << k << " must equal n_rows=" << n_rows; + randlapack_require(n == n_cols) << "VStackOp sketch: n=" << n << " must equal n_cols=" << n_cols; + + const int64_t d = m; // sketch output dimension + const int64_t b_blk = (block_size > 0) ? std::min(block_size, n) : std::min(256, n); + T* eye = new T[n_cols * b_blk](); + T* W = new T[n_rows * b_blk]; + for (int64_t j = 0; j < n; j += b_blk) { + int64_t b = std::min(b_blk, n - j); + std::fill_n(eye, n_cols * b, (T)0); + for (int64_t i = 0; i < b; ++i) eye[(j + i) + i * n_cols] = (T)1; + // W = [Top; Bot] * I_block (n_rows x b), via the dense NoTrans path. + (*this)(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, + n_rows, b, n_cols, (T)1, eye, n_cols, (T)0, W, n_rows); + // C[:, j:j+b] := alpha * op(S) * W + beta * C[:, j:j+b] + RandBLAS::sketch_general(Layout::ColMajor, trans_S, Op::NoTrans, + d, b, n_rows, alpha, S, W, n_rows, beta, C + j * ldc, ldc); + } + delete[] eye; + delete[] W; + } +}; + +} // namespace RandLAPACK::linops diff --git a/RandLAPACK/misc/rl_blas2_threads.hh b/RandLAPACK/misc/rl_blas2_threads.hh new file mode 100644 index 000000000..01041b9a7 --- /dev/null +++ b/RandLAPACK/misc/rl_blas2_threads.hh @@ -0,0 +1,230 @@ +#pragma once + +// Blas2ThreadGuard: caps the thread count of small dense LEVEL-2 BLAS calls +// (the triangular solves against an n x n preconditioner) for the duration of a +// scope, restoring the caller's setting on exit. +// +// WHY THIS EXISTS. +// +// A triangular solve with a single right-hand side is O(n^2) work on O(n^2) data: +// memory-bound, and with a sequential dependency chain. Threaded implementations +// pay a barrier per column (or per block), so barrier cost grows linearly in n +// while the useful work between barriers stays flat. Measured on a dense n = 2000 +// upper-triangular factor: +// +// threads 1 4 16 +// MKL dtrsv 0.448 ms 0.162 ms 31.3 ms +// +// This is not specific to triangularity (a plain dgemv of the same size degrades +// the same way) and does not improve with problem size (there is no crossover to +// wait for), so it is a fixed cost of threading a memory-bound, sequentially +// dependent kernel, not something blocking or a larger problem can fix. +// +// CONSEQUENCE IF LEFT UNGUARDED. The preconditioned least-squares solvers apply the +// preconditioner every inner iteration; the unpreconditioned baseline does not. The +// overhead therefore taxes exactly the methods that converge in few iterations, and +// can invert the wall-clock ranking against the unpreconditioned baseline by two +// orders of magnitude. +// +// SCOPE. Wrap ONLY the level-2 solves. The operator applies around them (FFTs, +// sparse solves, GEMMs) are level-3-like and genuinely want every thread, so a guard +// spanning a whole apply would trade one pathology for another. +// +// PORTABILITY. MKL is the only BLAS we can cap through a documented per-thread API, +// so the guard compiles to a no-op elsewhere. That is deliberate: other vendors are +// not known to thread trsv this way, and a global fallback (omp_set_num_threads) +// would leak the cap into concurrently running regions. +// +// TUNING. The cap defaults to kDefaultBlas2Threads and is overridable at runtime via +// the RANDLAPACK_BLAS2_THREADS environment variable (read once), so a new machine can +// be calibrated without a rebuild. A value <= 0 disables the guard entirely. + +#include + +#if defined(RandBLAS_HAS_MKL) +// mkl_service.h ONLY, never the umbrella : the latter redeclares the LAPACK +// entry points with MKL's own integer width, which conflicts with the declarations +// LAPACK++ already provides (measured: dozens of "conflicting declaration of C +// function" errors on ILP64 builds). The service header carries the threading +// controls and no BLAS/LAPACK prototypes. +#include +#endif + + +namespace RandLAPACK { + + +/// Size-dependent cap, calibrated on the benchmark hardware (Xeon Gold 6430, 64 +/// cores; dtrsv, milliseconds): +/// +/// threads 1 4 8 16 32 64 +/// n = 2000 0.562 0.408 *0.392* 0.774 0.909 0.887 +/// n = 8256 18.147 6.232 4.915 * 4.735* 5.828 7.130 +/// n = 20000 108.209 45.121 27.110 *19.276* 19.564 24.354 +/// +/// Threading genuinely helps up to 8-16 threads and degrades past that, so the +/// cap is a peak-seeker, not a "run it serially" switch. The optimum moves with +/// n because the parallel work per barrier grows with n while barrier cost does +/// not, hence the two-tier rule below. +/// +/// NOTE ON MAGNITUDE. On this hardware the penalty for leaving it unguarded (64 +/// threads) is a moderate 1.5-2.3x. A 16-thread WSL2 desktop showed a 100x +/// collapse instead, because 16 OpenMP threads there oversubscribe 8 physical +/// cores; do not quote desktop numbers as if they were cluster numbers. +constexpr int kBlas2ThreadsSmall = 8; ///< n <= kBlas2SmallDim +constexpr int kBlas2ThreadsLarge = 16; ///< n > kBlas2SmallDim +constexpr int64_t kBlas2SmallDim = 4000; + + +/// Thread cap for a level-2 solve on an n x n factor. RANDLAPACK_BLAS2_THREADS +/// overrides the calibrated values (read once); a value <= 0 disables the guard. +inline int blas2_thread_cap(int64_t n) { + static const int override_cap = []() -> int { + const char* s = std::getenv("RANDLAPACK_BLAS2_THREADS"); + if (s == nullptr || *s == '\0') return -1; // -1 = "no override" + return std::atoi(s); + }(); + if (override_cap >= 0) return override_cap; + return (n <= kBlas2SmallDim) ? kBlas2ThreadsSmall : kBlas2ThreadsLarge; +} + + +/// Cap for MKL's threaded FFT (DFTI). Measured on the benchmark node (Xeon Gold +/// 6430, 64 cores) for one forward+backward apply of the Toeplitz operator, +/// milliseconds: +/// +/// threads 1 8 16 32 64 +/// L = 32768 0.274 0.355 0.267 0.297 0.319 +/// L = 131072 3.29 1.24 1.00 0.93 0.97 +/// L = 524288 15.52 4.95 3.66 3.05 3.13 +/// +/// The mean barely improves past 16 threads, and at 64 the call becomes +/// INTERMITTENTLY unstable (individual transforms occasionally take 100x +/// their typical cost), which a solver that converges in a handful of +/// iterations cannot average out. Capping the transform is the +/// vendor-recommended remedy for single small transforms (Intel advises +/// reducing the thread count rather than expecting a lone FFT to scale). +/// Batching via DFTI_NUMBER_OF_TRANSFORMS is the remedy for the multi-column +/// build path (sketch and Gram applies); the CG loop still produces one +/// right-hand side at a time and keeps this cap, further narrowed by +/// SolveWidthScope while a solver is running. +/// +/// MKL decides an FFT descriptor's threading at commit time in general, but a +/// cap applied only around DftiComputeForward/Backward (via Blas2ThreadGuard +/// below) was measured to still bind FFT width on this hardware even for a +/// descriptor committed at ambient width. ext_toeplitz_linop.hh also sets +/// DFTI_THREAD_LIMIT at commit time; the two caps compose, and the narrower +/// compute-time cap wins. +constexpr int kDefaultFFTThreads = 16; + + +/// Thread cap for an FFT apply. RANDLAPACK_FFT_THREADS overrides (read once); +/// <= 0 disables the guard. +inline int fft_thread_cap() { + static const int cap = []() -> int { + const char* s = std::getenv("RANDLAPACK_FFT_THREADS"); + if (s == nullptr || *s == '\0') return kDefaultFFTThreads; + return std::atoi(s); + }(); + return cap; +} + + +/// Solve-scoped width matching. Alternating OpenMP team widths cost the wider +/// region ~300 us per re-formation on the benchmark node (libgomp, dual-socket +/// Gold 6430). Inside an iterative solve the trsv runs at blas2_thread_cap(n) +/// ACTUAL width, and widths can only be equalized DOWNWARD: MKL does not form +/// wide teams for small trsv, so raising the trsv request does not remove the +/// alternation. The scope below therefore narrows every width-capped kernel +/// that consults it (currently the Toeplitz FFT apply) to the trsv width for +/// the duration of a solver call, so the inner loop runs at ONE width. +/// Build-phase applies see no active scope and keep their own calibrated caps. +/// RANDLAPACK_SOLVE_FFT_MATCH=0 disables the matching (read once; for A/B probes). +inline bool solve_width_match_enabled() { + static const bool on = []() { + const char* s = std::getenv("RANDLAPACK_SOLVE_FFT_MATCH"); + return !(s != nullptr && s[0] == '0' && s[1] == '\0'); + }(); + return on; +} + +/// The active solve-context width for the calling thread; 0 = no active scope. +inline int& solve_context_width_ref() { + thread_local int width = 0; + return width; +} +inline int solve_context_width() { return solve_context_width_ref(); } + +/// RAII solve-width context. Instantiated by the iterative solvers for their +/// whole duration; width-capped kernels take min(own cap, context width) while +/// one is active. Nesting restores the enclosing context on destruction. +class SolveWidthScope { + public: + /// @param n preconditioner dimension; the context width is the trsv cap + /// blas2_thread_cap(n), the narrow width the loop already pays. + explicit SolveWidthScope(int64_t n) { + const int cap = blas2_thread_cap(n); + if (solve_width_match_enabled() && cap > 0) { + prev_ = solve_context_width_ref(); + solve_context_width_ref() = cap; + active_ = true; + } + } + ~SolveWidthScope() { + if (active_) solve_context_width_ref() = prev_; + } + SolveWidthScope(const SolveWidthScope&) = delete; + SolveWidthScope& operator=(const SolveWidthScope&) = delete; + SolveWidthScope(SolveWidthScope&&) = delete; + SolveWidthScope& operator=(SolveWidthScope&&) = delete; + private: + int prev_ = 0; + bool active_ = false; +}; + + +/// RAII cap on the calling thread's MKL thread count. Construct in the narrowest +/// scope containing the guarded call; the previous setting is restored on +/// destruction (including when an exception unwinds through the scope). +class Blas2ThreadGuard { + public: + /// @param n dimension of the triangular factor being solved against; + /// selects the calibrated cap (see blas2_thread_cap). + explicit Blas2ThreadGuard(int64_t n) : Blas2ThreadGuard(blas2_thread_cap(n), 0) {} + + /// Explicit-cap form, for callers with their own calibration (e.g. the FFT + /// operator). The dummy second parameter disambiguates from the int64_t + /// dimension overload. + Blas2ThreadGuard(int cap, int /*tag*/) { + #if defined(RandBLAS_HAS_MKL) + if (cap > 0) { + // mkl_set_num_threads_local returns the PREVIOUS thread-local value, + // where 0 means "no local setting, follow the global one". Restoring + // that value in the destructor therefore also restores the + // follow-the-global state, rather than pinning the global count. + prev_ = mkl_set_num_threads_local(cap); + active_ = true; + } + #endif + } + + ~Blas2ThreadGuard() { + #if defined(RandBLAS_HAS_MKL) + if (active_) mkl_set_num_threads_local(prev_); + #endif + } + + Blas2ThreadGuard(const Blas2ThreadGuard&) = delete; + Blas2ThreadGuard& operator=(const Blas2ThreadGuard&) = delete; + Blas2ThreadGuard(Blas2ThreadGuard&&) = delete; + Blas2ThreadGuard& operator=(Blas2ThreadGuard&&) = delete; + + private: + #if defined(RandBLAS_HAS_MKL) + int prev_ = 0; + #endif + bool active_ = false; +}; + + +} // namespace RandLAPACK diff --git a/RandLAPACK/testing/rl_gen.hh b/RandLAPACK/testing/rl_gen.hh index abf6cf812..81b8e62bf 100644 --- a/RandLAPACK/testing/rl_gen.hh +++ b/RandLAPACK/testing/rl_gen.hh @@ -358,22 +358,41 @@ void gen_oleg_adversarial_mat( } /// Generate singular values for the "bad CholQR" matrix. -/// The first k values are 1, then values start at 10^-8 and decrease -/// exponentially, controlled by cond and n. /// -/// @param[in] k Number of singular values (= sketching dimension) -/// @param[in] n Number of columns in the target matrix -/// @param[in] cond Condition number +/// The leading half of the spectrum is set to one. The trailing half drops to +/// 1e-8 and then decays geometrically to 1/cond, so the returned spectrum has +/// condition number exactly cond. The cliff between the two blocks is the point +/// of this input: it is what drives the Gram matrix numerically indefinite, and +/// so exposes the failure mode of an unshifted CholeskyQR. /// -/// @return Vector of k singular values +/// Requires cond >= 1e8. Below that threshold the trailing block would rise from +/// 1e-8 toward 1/cond rather than decay, leaving a non-monotone spectrum whose +/// condition number is 1e8 rather than the requested value. That threshold is +/// also where the failure being modeled begins, since an unshifted CholeskyQR +/// loses orthogonality once cond exceeds eps^(-1/2), about 1.5e8 in double. +/// +/// The previous version of this routine took an unused second dimension argument +/// and computed an empty loop, returning all ones (condition number 1) for every +/// requested cond. It had no callers other than gen_bad_cholqr_mat below. +/// +/// @param[in] k Number of singular values to generate. Must be >= 2. +/// @param[in] cond Target condition number. Must be >= 1e8. +/// +/// @return Vector of k singular values, non-increasing, with s[0] = 1 and +/// s[k-1] = 1/cond. template -std::vector gen_bad_cholqr_singvals(int64_t k, int64_t n, T cond) { +std::vector gen_bad_cholqr_singvals(int64_t k, T cond) { + randlapack_require(k >= 2) << "k=" << k << " must be >= 2 to admit both a leading and a trailing block"; + randlapack_require(cond >= T(1e8)) << "cond=" << cond << " must be >= 1e8; below that the trailing block is not monotone"; + std::vector s(k, 1.0); - int offset = k; - T t = log(std::pow(10, 8) / cond) / (1 - (n - offset)); - T cnt = 0.0; - for (int i = offset; i < k; ++i) { - s[i] = (std::exp(t) / std::pow(10, 8)) * (std::exp(++cnt * -t)); + int64_t offset = k / 2; // size of the leading block of ones + int64_t n_decay = k - offset; // size of the trailing block + + // Geometric interpolation from 1e-8 down to 1/cond across the trailing block. + for (int64_t i = 0; i < n_decay; ++i) { + T frac = (n_decay == 1) ? T(0) : T(i) / T(n_decay - 1); + s[offset + i] = T(1e-8) * std::pow(cond * T(1e-8), -frac); } return s; } @@ -390,7 +409,7 @@ void gen_bad_cholqr_mat( bool diagon, RandBLAS::RNGState &state ) { - auto s = gen_bad_cholqr_singvals(k, n, cond); + auto s = gen_bad_cholqr_singvals(k, cond); T* S = new T[k * k](); RandLAPACK::util::diag(k, k, s.data(), k, S); @@ -514,6 +533,37 @@ void gen_random_dense( } } +/// Generate a symmetric tridiagonal matrix in RandBLAS CSR format (no Eigen). +/// +/// Row i carries `offdiag` at columns i-1 and i+1 (within bounds) and `diag` on +/// the diagonal, with column indices stored in ascending order per row. With +/// diag = 2, offdiag = -1 this is the SPD 1D Laplacian (a convenient PD testbed +/// for sparse solvers); other (diag, offdiag) give indefinite or non-diagonally- +/// dominant variants. +/// +/// @tparam T Scalar type. +/// @tparam sint_t CSR index type (default int64_t). +/// @param[in] n Matrix dimension (n x n), n >= 1. +/// @param[in] diag Diagonal value. +/// @param[in] offdiag Sub/super-diagonal value (symmetric). +/// @return A newly-owned CSR matrix with nnz = 3n - 2 (n >= 2) or 1 (n == 1). +template +RandBLAS::sparse_data::CSRMatrix gen_tridiag_csr(int64_t n, T diag, T offdiag) { + randblas_require(n >= 1); + int64_t nnz = (n == 1) ? 1 : 3 * n - 2; + RandBLAS::sparse_data::CSRMatrix A(n, n); + A.reserve(nnz); + int64_t p = 0; + for (int64_t i = 0; i < n; ++i) { + A.rowptr[i] = static_cast(p); + if (i > 0) { A.colidxs[p] = static_cast(i - 1); A.vals[p] = offdiag; ++p; } + { A.colidxs[p] = static_cast(i); A.vals[p] = diag; ++p; } + if (i + 1 < n) { A.colidxs[p] = static_cast(i + 1); A.vals[p] = offdiag; ++p; } + } + A.rowptr[n] = static_cast(p); + return A; +} + /// Generate a random sparse matrix in COO format. /// Creates a sparse matrix with uniformly random positions and values from the specified distribution. /// Duplicate (row, col) entries are merged by summing their values. diff --git a/RandLAPACK/testing/rl_memory_tracker.hh b/RandLAPACK/testing/rl_memory_tracker.hh index 30a80c531..39ef7350a 100644 --- a/RandLAPACK/testing/rl_memory_tracker.hh +++ b/RandLAPACK/testing/rl_memory_tracker.hh @@ -4,6 +4,8 @@ // - Peak RSS sampling via background thread // - Analytical peak working memory computation for each algorithm +#include "rl_exceptions.hh" + #include #include #include @@ -12,6 +14,10 @@ #include #include #include +#include +#if defined(__GLIBC__) || defined(__linux__) +#include // malloc_trim +#endif namespace RandLAPACK { @@ -40,7 +46,35 @@ static inline long get_rss_kb() { // long peak_increase_kb = tracker.stop(); class PeakRSSTracker { public: + // Joins a still-running sampler thread instead of letting a joinable + // std::thread destruct (which calls std::terminate). Reached when start() + // was called but stop() was skipped, e.g. an exception thrown in between. + ~PeakRSSTracker() { + if (sampler_.joinable()) { + running_.store(false, std::memory_order_relaxed); + sampler_.join(); + } + } + void start() { + // A second start() without an intervening stop() would assign into an + // already-joinable sampler_ and call std::terminate; require stop() first. + randlapack_require(!sampler_.joinable()) + << "PeakRSSTracker::start: sampler already running; call stop() first"; + + // Release freed heap back to the OS before taking the baseline. + // RSS is process-cumulative: glibc keeps freed arenas + // mapped, so without this the FIRST tracked algorithm in a benchmark + // absorbs the whole process ramp-up (its delta over-reports) while + // every later one reuses already-faulted pages (delta ~0, the + // "peak_rss_kb=4" effect). Trimming resets + // the floor to live memory only, making per-method deltas comparable + // regardless of execution order. Frees only unused arena space; no + // effect on correctness or on MKL's internal buffers (a warmup pass + // should absorb those). +#if defined(__GLIBC__) + malloc_trim(0); +#endif baseline_kb_ = get_rss_kb(); peak_kb_.store(baseline_kb_, std::memory_order_relaxed); running_.store(true, std::memory_order_relaxed); @@ -81,50 +115,165 @@ private: // These compute the peak memory from known buffer sizes in each algorithm, // excluding test-mode Q-factor allocation (which is only for verification). // All return memory in KB. +// +// SCOPE. These model the DRIVER's workspace only. They do NOT model: +// * the operator's per-apply temporaries (CompositeOperator allocates a +// fresh scratch buffer on every apply, once per nesting level; negligible +// at small block sizes, but can dominate at b_eff = n). +// * the operator's SKETCH path (VStackOp's sketch overload allocates +// scratch proportional to (m + n) * b_blk; every sketching method pays +// it, CholQR/CholQR2 do not sketch). This is deliberately NOT folded into +// the formulas below: it is a property of the operator passed in, not of +// the driver, and would be wrong for an operator whose sketch path +// allocates differently. +// * operator state allocated once and reused (e.g. the Toeplitz FFT plans +// and batch buffer), which the benchmark warmup faults in before the +// tracker baseline is taken, so it cancels out of the per-row delta. +// Consequence: peak-vs-predicted is a meaningful check for blocked methods on +// any operator, and for non-blocked methods only on operators whose applies +// allocate nothing (e.g. the matrix-free Toeplitz operator). +// +// PARAMETER CONVENTION: `m` is the ROW COUNT OF THE OPERATOR THE QR RUNS ON. +// For an augmented operator A_hat = [A; mu I] that is m + n, not m. +// +// `d` matches the drivers' truncating cast, not ceil. // --------------------------------------------------------------------------- -// CQRRT_linops: A_hat(d*n) + tau(n) + R_sk_inv(n*n) + A_pre(m*b_eff) +// CQRRT_linops (TRSM_IDENTITY / GEQP3), PHASED allocation. +// The sketch phase and the Gram/Cholesky phase have disjoint working sets, and the +// driver now allocates each only while it is needed, so the peak is the MAX of the +// two moments rather than their sum: +// sketch moment : A_hat(d*n) + tau(n) + P(n*n) +// Gram moment : P + R_pre + G (3*n*n) + A_temp(m*b_eff) + Z_buf(n*b_eff) +// + diag_backup(n), allocated inside cholqr_primitive whenever +// retries are enabled (the driver default: max_retries = -1). +// cholqr_primitive's shift-retry no longer keeps a full n x n Gram backup: a failed +// potrf attempt is undone from the Gram's own untouched strict lower triangle plus +// an O(n) diagonal snapshot (see rl_cholqr.hh), so the retry scratch is +n, not +n*n. +// For any d <= 3n the Gram moment dominates, which puts CQRRT at exactly the +// CholQR2 / sCholQR3 peak. Phasing (splitting the sketch and Gram allocations so +// they don't coexist) is still a real cut versus their sum, d*n + n + 3*n*n + +// (m+n)*b_eff + n, for any d > n. template static inline long cqrrt_linops_analytical_kb(int64_t m, int64_t n, double d_factor, int64_t block_size) { - int64_t d = static_cast(std::ceil(d_factor * n)); + int64_t d = static_cast(d_factor * n); if (d < n) d = n; // matches the drivers int64_t b_eff = (block_size > 0 && block_size < n) ? block_size : n; - long bytes = static_cast(sizeof(T)) * (d * n + n + (long)n * n + (long)m * b_eff); - return bytes / 1024; + long sketch_moment = static_cast(sizeof(T)) * ((long)d * n + n + (long)n * n); + long gram_moment = static_cast(sizeof(T)) * (3L * n * n + (long)(m + n) * b_eff + n); + return std::max(sketch_moment, gram_moment) / 1024; } -// CholQR_linops: I_mat(n*n) + A_temp(m*b_eff) +// CholQR_linops (adaptive-shift retries enabled by default): +// cholqr_primitive owns G(n*n) + A_temp(m*b_eff) + diag_backup(n) (the O(n) +// retry snapshot; allocated whenever retries are enabled, the driver default) +// plus blocked_preconditioned_gram's I_block(n*b_eff) inside the Gram loop. +// Peak = n*n + (m+n)*b_eff + n. template static inline long cholqr_linops_analytical_kb(int64_t m, int64_t n, int64_t block_size = 0) { int64_t b_eff = (block_size > 0 && block_size < n) ? block_size : n; - long bytes = static_cast(sizeof(T)) * ((long)n * n + (long)m * b_eff); + long bytes = static_cast(sizeof(T)) * (1L * n * n + (long)(m + n) * b_eff + n); return bytes / 1024; } -// sCholQR3_linops (fully-blocked): G(n*n) + R_temp(n*n) + M(n*n) + A_temp(m*b) + Z_buf(n*b) -// No m x n buffer during QR iterations; all Gram matrices computed through blocked linop calls. +// CholQR2_linops: call() is a thin wrapper around cholqr_iterate(num_iters=2) +// (comps/rl_cholqr.hh), which owns the scratch directly (no per-class member +// buffers). Iter 1 allocates and frees its own G(n*n) + A_temp(m*b_eff) inside +// cholqr_primitive's unpreconditioned overload BEFORE the iter-2 scratch below +// is allocated, so the two moments are sequential, not coexistent; iter 1 is +// bounded by the iter-2 peak. Peak (iter 2, preconditioned): +// cholqr_iterate's persistent G, R_pre, P_prev (3 n^2) + A_temp (m * b_eff) +// + Z_buf (n * b_eff) + cholqr_primitive's diag_backup(n), the O(n) retry +// snapshot allocated whenever retries are enabled (the driver default). +// See the SCOPE note at the top of this section. +template +static inline long cholqr2_linops_analytical_kb(int64_t m, int64_t n, int64_t block_size) { + int64_t b_eff = (block_size > 0 && block_size < n) ? block_size : n; + long bytes = static_cast(sizeof(T)) * + ( 3 * n * n // G + R_pre + P_prev (iter-2 persistent scratch) + + m * b_eff // A_temp + + n * b_eff // Z_buf + + n // diag_backup (retry snapshot) + ); + return bytes / 1024; +} + +// sCholQR3_linops: call() is cholqr_iterate(num_iters=3) (comps/rl_cholqr.hh); +// same sequencing argument as CholQR2_linops above (iter 1's primitive frees its +// own G/A_temp BEFORE the iters-2/3 scratch is allocated, so iter 1 is bounded +// by the iters-2/3 peak). +// Peak moment (iters 2/3): driver persistent G(n*n) + R_pre(n*n) + P_prev(n*n) +// + A_temp(m*b_eff) + Z_buf(n*b_eff) +// + primitive diag_backup(n), the O(n) retry snapshot (rl_cholqr.hh: a failed +// potrf attempt is restored from G's own untouched strict lower triangle +// plus this diagonal snapshot, not a full n x n Gram backup). +// = 3*n*n + (m+n)*b_eff + n, which also bounds the iter-1 moment. template static inline long scholqr3_linops_analytical_kb(int64_t m, int64_t n, int64_t block_size = 0) { int64_t b_eff = (block_size > 0 && block_size < n) ? block_size : n; - long bytes = static_cast(sizeof(T)) * (3L * n * n + (long)(m + n) * b_eff); + long bytes = static_cast(sizeof(T)) * (3L * n * n + (long)(m + n) * b_eff + n); return bytes / 1024; } -// sCholQR3_linops_basic: Q_buf(m*n) + G(n*n) + R_temp(n*n) + M(n*n) -// Materializes Q = A * R1^{-1} after iteration 1, then uses dense syrk for iterations 2-3. -// No blocking — always O(m*n + n^2) peak. +// sCholQR3_linops_basic: same sequencing argument as the blocked variant; b_eff = n +// collapse of 3*n*n + (m+n)*b_eff + n gives 3*n*n + (m+n)*n + n = 4*n*n + m*n + n. +// Validated against the formula above on the matrix-free Toeplitz operator: +// ratio 0.987 to 1.028 (n = 1000 and n = 4000), superseding the earlier +// validation against the pre-diag-backup 5*n*n + m*n formula. On an operator whose +// applies allocate (FEM2's nested CompositeOperator) this driver-only figure is +// NOT the process peak: b_eff = n makes each per-apply temporary inner_dim * n, +// and FEM2 large measured 255 GB against a 126 GB driver workspace. See the +// SCOPE note at the top of this section before quoting this number. template static inline long scholqr3_linops_basic_analytical_kb(int64_t m, int64_t n) { - long bytes = static_cast(sizeof(T)) * ((long)m * n + 3L * n * n); + long bytes = static_cast(sizeof(T)) * (4L * n * n + (long)m * n + n); + return bytes / 1024; +} + +// Blendenpik_linops (sketch + Householder QR + LSQR): R is handed to the caller +// as R_out (ownership transfer, rl_blendenpik.hh: `R_out = R; R = nullptr;`) with +// no copy, so R_out never coexists with a separate R buffer: R_out just IS R. +// Buffers, all live simultaneously at the peak because call() holds the sketch +// (Ask, tau) allocated until cleanup() at the very end, well after R is done +// with it: +// Ask(d*n) + tau(n) + R(n*n) +// ws (warm_start || init_only) only : Sb(d) + x0(n) + r0(m) +// LSQR workspace (with_lsqr only) : u(m) + av(m) + v(n) + w(n) + atu(n) +// + sc(n) = 2m + 4n +// init_only (with_lsqr = false, the refined rows' mode) forces ws = true inside +// call() regardless of the warm_start member: init_only implies ws by +// construction, so the x0-build buffers are always live in that mode. The +// `warm_start` parameter here therefore only gates the vectors when with_lsqr is +// true (the published Blendenpik/Blendenpik_cold rows, where ws genuinely +// depends on the member); when with_lsqr is false it is ignored and treated as +// true, so a caller that mislabels a cold init_only row still gets a faithful +// prediction. init_only rows then continue into restarted_pcg_ne, whose own +// workspace (9n + m) is smaller than the sketch term for any d >= 1, so the +// Blendenpik moment remains the peak. +template +static inline long blendenpik_linops_analytical_kb(int64_t m, int64_t n, double d_factor, + bool warm_start = true, bool with_lsqr = true) { + int64_t d = static_cast(d_factor * n); if (d < n) d = n; // matches the driver + bool ws = with_lsqr ? warm_start : true; // init_only always forces ws = true + long vecs = ws ? ((long)d + n + m) : 0L; // Sb, x0, r0 + if (with_lsqr) vecs += 2L * m + 4L * n; // u, av | v, w, atu, sc + long bytes = static_cast(sizeof(T)) * ((long)d * n + n + (long)n * n + vecs); return bytes / 1024; } // Dense CQRRT (materialize + rl_cqrrt): -// Peak = A_materialized(m*n) + A_hat(d*n) + tau(n) -// (I_mat freed before rl_cqrrt allocates A_hat, so they don't overlap) +// Peak = A_materialized(m*n) + A_hat(d*n) + tau(n) + gram_backup(n*n) +// (I_mat freed before rl_cqrrt allocates A_hat, so they don't overlap; A_hat +// stays live through the whole call, so it coexists with gram_backup). +// gram_backup is allocated whenever rl_cqrrt's max_retries != 0 (the driver +// default is -1, unbounded retries, so it is live in the common case); +// retries_enabled defaults to true to match that default and the current +// CQRRT_linop_basic.cc caller, which never overrides max_retries. template -static inline long dense_cqrrt_analytical_kb(int64_t m, int64_t n, double d_factor) { - int64_t d = static_cast(std::ceil(d_factor * n)); - long bytes = static_cast(sizeof(T)) * ((long)m * n + (long)d * n + n); +static inline long dense_cqrrt_analytical_kb(int64_t m, int64_t n, double d_factor, + bool retries_enabled = true) { + int64_t d = static_cast(d_factor * n); if (d < n) d = n; // matches the drivers + long gram_backup = retries_enabled ? (long)n * n : 0L; + long bytes = static_cast(sizeof(T)) * ((long)m * n + (long)d * n + n + gram_backup); return bytes / 1024; } diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt index 6aa7749d1..a2e03957f 100644 --- a/benchmark/CMakeLists.txt +++ b/benchmark/CMakeLists.txt @@ -156,8 +156,89 @@ set( Eigen3::Eigen fast_matrix_market ) -add_benchmark(NAME CQRRT_linop_composite_applications CXX_SOURCES bench_CQRRT_linops/CQRRT_linop_composite_applications.cc LINK_LIBS ${Benchmark_libs_cqrrt_linops}) +add_benchmark(NAME CQRRT_linop_applications CXX_SOURCES bench_CQRRT_linops/CQRRT_linop_applications.cc LINK_LIBS ${Benchmark_libs_cqrrt_linops}) add_benchmark(NAME CQRRT_linop_basic CXX_SOURCES bench_CQRRT_linops/CQRRT_linop_basic.cc LINK_LIBS ${Benchmark_libs_cqrrt_linops}) +add_benchmark(NAME CQRRT_diagnostic CXX_SOURCES bench_CQRRT_linops/CQRRT_diagnostic.cc LINK_LIBS ${Benchmark_libs_cqrrt_linops}) + +# Toeplitz LS benchmark (autoregression/system-id experiment). The matrix-free +# prolate-Toeplitz operator (extras/linops/ext_toeplitz_linop.hh) uses MKL's FFT (DFTI), +# so this target additionally needs the MKL include dir + the ILP64 define (the MKL libs +# themselves are already linked transitively via RandLAPACK -> blaspp -> MKL). RandLAPACK +# also supports OpenBLAS, so on a non-MKL machine mkl_dfti.h will not exist; skip only +# this target with a warning rather than aborting the whole benchmark configure. +# Locate mkl_dfti.h. MKLROOT is the usual signal, but batch builds (e.g. ISAAC's +# fresh-clone install job) may not export it, so also probe the other oneAPI env +# vars and common install roots. As a last resort, derive it from a BLAS/LAPACK +# library dir that CMake already resolved (MKL's lib and include are siblings). +find_path(MKL_DFTI_INCLUDE_DIR mkl_dfti.h + HINTS + $ENV{MKLROOT}/include + $ENV{MKL_ROOT}/include + $ENV{ONEAPI_ROOT}/mkl/latest/include + $ENV{CMPLR_ROOT}/../mkl/latest/include + /opt/intel/oneapi/mkl/latest/include + /opt/intel/oneapi/mkl/2026.0/include + PATH_SUFFIXES include) +if(NOT MKL_DFTI_INCLUDE_DIR) + # Derive from a resolved MKL/LAPACK library path (…/mkl//lib -> …/include). + foreach(_lib ${BLAS_LIBRARIES} ${LAPACK_LIBRARIES} ${blaspp_libraries}) + if(_lib MATCHES "mkl") + get_filename_component(_libdir "${_lib}" DIRECTORY) + find_path(MKL_DFTI_INCLUDE_DIR mkl_dfti.h HINTS "${_libdir}/../include" "${_libdir}/../../include") + endif() + endforeach() +endif() +if(NOT MKL_DFTI_INCLUDE_DIR) + message(WARNING "mkl_dfti.h not found; skipping toeplitz_ls_benchmark target. Pass " + "-DMKL_DFTI_INCLUDE_DIR=/include or export MKLROOT to build it.") +else() + message(STATUS "Toeplitz benchmark: mkl_dfti.h in ${MKL_DFTI_INCLUDE_DIR}") + add_benchmark(NAME toeplitz_ls_benchmark CXX_SOURCES bench_toeplitz_ls/toeplitz_ls_benchmark.cc LINK_LIBS ${Benchmark_libs_cqrrt_linops}) + target_include_directories(toeplitz_ls_benchmark PRIVATE ${MKL_DFTI_INCLUDE_DIR}) + + # Derive MKL_ILP64 from the BLAS interface blaspp was actually built against, + # instead of assuming ILP64 (an LP64 MKL build would otherwise silently mismatch + # the DFTI header's ABI). blaspp's installed config records this in + # blas/defines.h (BLAS_ILP64 defined iff the ILP64 interface was linked); + # blaspp_DIR locates that install tree. Fall back to matching ilp64/lp64 in the + # resolved library names if that header cannot be found. + set(_toeplitz_ilp64 "") + if(blaspp_DIR) + get_filename_component(_blaspp_prefix "${blaspp_DIR}/../../.." ABSOLUTE) + set(_blaspp_defines_h "${_blaspp_prefix}/include/blas/defines.h") + if(EXISTS "${_blaspp_defines_h}") + file(READ "${_blaspp_defines_h}" _blaspp_defines_contents) + if(_blaspp_defines_contents MATCHES "#define[ \t]+BLAS_ILP64") + set(_toeplitz_ilp64 TRUE) + else() + set(_toeplitz_ilp64 FALSE) + endif() + message(STATUS "Toeplitz benchmark: BLAS interface from ${_blaspp_defines_h}: " + "ILP64=${_toeplitz_ilp64}") + endif() + endif() + if(_toeplitz_ilp64 STREQUAL "") + foreach(_lib ${BLAS_LIBRARIES} ${LAPACK_LIBRARIES} ${blaspp_libraries}) + if(_lib MATCHES "ilp64") + set(_toeplitz_ilp64 TRUE) + elseif(_lib MATCHES "lp64") + set(_toeplitz_ilp64 FALSE) + endif() + endforeach() + if(NOT _toeplitz_ilp64 STREQUAL "") + message(STATUS "Toeplitz benchmark: BLAS interface from linked library names: " + "ILP64=${_toeplitz_ilp64}") + endif() + endif() + if(_toeplitz_ilp64 STREQUAL "") + message(WARNING "Toeplitz benchmark: could not determine whether the resolved " + "BLAS/MKL interface is ILP64 or LP64; defaulting to MKL_ILP64.") + set(_toeplitz_ilp64 TRUE) + endif() + if(_toeplitz_ilp64) + target_compile_definitions(toeplitz_ls_benchmark PRIVATE MKL_ILP64) + endif() +endif() # ABRIK benchmarks add_benchmark(NAME ABRIK_runtime_breakdown CXX_SOURCES bench_ABRIK/ABRIK_runtime_breakdown.cc LINK_LIBS ${Benchmark_libs}) diff --git a/benchmark/bench_CQRRT_linops/CQRRT_diagnostic.cc b/benchmark/bench_CQRRT_linops/CQRRT_diagnostic.cc new file mode 100644 index 000000000..019773d66 --- /dev/null +++ b/benchmark/bench_CQRRT_linops/CQRRT_diagnostic.cc @@ -0,0 +1,785 @@ +// CQRRT preconditioner comparison benchmark +// +// Isolates the effect of different methods for forming R_sk^{-1} on the final +// orthogonality quality of CQRRT. Tests seven paths: +// +// [1] expl_trsm: DTRSM_R(A, R_sk) in-place <- CQRRT_expl path +// [2] expl_inv_trsm_left: solve R_sk * X = I (TRSM Side::Left) -> R_inv; DGEMM(A, R_inv) +// (column-ordered solve; the RandLAPACK default, +// PCholQRPrecondMethod::TRSM_IDENTITY) +// [3] expl_inv_trsm_right: solve X * R_sk = I (TRSM Side::Right) -> R_inv; DGEMM(A, R_inv) +// (reversed, row-ordered solve; kept to document the +// solve-ordering effect, NOT shipped) +// [4] expl_inv_trtri: TRTRI(R_sk) -> R_inv; DGEMM(A, R_inv) (LAPACK trtri) +// [5] expl_inv_geqp3: GEQP3(R_sk) = Q*R_buf*P^T; +// R_inv = P * TRSM(R_buf, Q^T); DGEMM(A, R_inv) +// [6] expl_inv_svd: GESDD(R_sk) = U*S*Vt; +// R_inv = V * diag(1/S) * U^T; DGEMM(A, R_inv) +// [7] expl_inv_bqrrp: BQRRP(R_sk)=Q*R_buf*P^T; R_inv=P*TRSM(R_buf,Q^T); DGEMM(A, R_inv) +// +// Path [1] never forms R_sk^{-1} explicitly (backward stable). +// Paths [2]-[7] all form R_sk^{-1} explicitly via different methods. +// Paths [2] and [3] differ ONLY in how the triangular system is posed; +// the ordering governs the error of the composite product M * R_inv. +// Path [5] uses a rank-revealing QR to invert R_sk; the Q factor makes +// the inversion well-conditioned even when R_sk itself is ill-conditioned. +// Path [6] uses the SVD (gold standard for stability). +// Path [7] uses BQRRP (blocked randomized QRCP), the randomized +// counterpart of path [5]'s GEQP3. +// +// All seven paths use the same sketch (same RNG state). +// +// Per-path metrics: +// cond(A_pre) : condition number of the preconditioned matrix +// cond(G = A_pre^T A_pre) : condition number of the Gram matrix (input to Cholesky) +// orth_error(Q) : full-pipeline: G=SYRK(A_pre), R_chol=chol(G), +// R_final=R_chol*R_sk, Q=A_orig*R_final^{-1} +// +// Cross-path relative differences of A_pre (reference = path [1]): +// rd_1p = ||A_pre[1] - A_pre[p]|| / ||A_pre[1]|| for p = 2..7 +// +// Step-by-step pipeline divergence between paths [1] and [2] (CQRRT_expl vs CQRRT_linop): +// Each path uses the same RNG seed but a different sketch code path, mirroring actual impls. +// Path [1] (CQRRT_expl): sketch via sketch_general(S, A_dense); TRSM in-place; SYRK +// Path [2] (CQRRT_linop): sketch via A_linop(Side::Right, S) [SpGEMM]; +// TRSM_IDENTITY → R_inv; A_linop fwd/adj; +// TRSM(Left,Trans) Gram completion on R_sk; TRMM R_final +// rd_Msk_12 = ||Ahat1 - Ahat2|| / ||Ahat1|| (raw sketch S*A, different code paths) +// rd_Rsk_12 = ||R_sk1 - R_sk2|| / ||R_sk1|| (QR of above) +// rd_G_12 = ||G1 - G2|| / ||G1|| (Gram matrix) +// rd_Rchol_12 = ||Rchol1 - Rchol2|| / ||Rchol1|| (Cholesky factor) +// rd_Rfinal_12= ||Rfinal1 - Rfinal2|| / ||Rfinal1|| (final R = R_chol * R_sk) +// +// Sketch diagnostic: +// cond(R_sk) +// +// Usage (file mode): +// ./CQRRT_diagnostic [sketch_nnz] +// +// Usage (generate mode): +// ./CQRRT_diagnostic gen [sketch_nnz] +// +// NOTE: fidelity gaps vs the shipped drivers (cholqr_primitive, comps/rl_cholqr.hh): +// - No adaptive Cholesky-shift retry is modeled here; a potrf failure in Part B +// is left unretried, so failures are expected in the high-kappa regime this +// tool probes. +// - The RANDLAPACK_GRAM_LEFT=gemm per-block-GEMM Gram completion arm is not +// modeled; this tool always completes the Gram with a single TRSM. +// - R_final is formed via TRMM(Side::Right, R_sk) here, versus the primitive's +// TRMM(Side::Left, R_chol); same mathematical product, different rounding. +// - This tool computes the sketch dimension d = ceil(d_factor * n), while the +// shipped drivers (rl_cqrrt.hh) truncate: d = (int64_t)(d_factor * n). The +// two agree only when d_factor * n is exactly representable; otherwise this +// tool's sketch is one row taller than the driver it is meant to model. + +#include "RandLAPACK.hh" +#include "rl_blaspp.hh" +#include "rl_lapackpp.hh" +#include "rl_gen.hh" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#ifdef _OPENMP +#include +#endif + +#include "../../extras/misc/ext_util.hh" +#include "RandLAPACK/testing/rl_test_utils.hh" +#include "cqrrt_bench_common.hh" + +using std::chrono::steady_clock; +using std::chrono::duration_cast; +using std::chrono::microseconds; +using blas::Layout; +using blas::Op; +using blas::Side; +using blas::Uplo; +using blas::Diag; +using RandLAPACK::bench::quote_join_argv; +using RandLAPACK::bench::get_hostname; + +// ============================================================================ +// Path constants +// ============================================================================ + +static constexpr int N_PATHS = 7; + +static constexpr const char* PATH_NAMES[N_PATHS] = { + "expl_trsm", + "expl_inv_trsm_left", + "expl_inv_trsm_right", + "expl_inv_trtri", + "expl_inv_geqp3", + "expl_inv_svd", + "expl_inv_bqrrp", +}; + +static constexpr const char* PATH_DESCS[N_PATHS] = { + "DTRSM_R(A, R_sk) in-place <- CQRRT_expl", + "solve R_sk*X=I (Side::Left)->R_inv; DGEMM(A, R_inv) <- RandLAPACK default", + "solve X*R_sk=I (Side::Right)->R_inv; DGEMM(A, R_inv) (reversed ordering, NOT shipped)", + "TRTRI(R_sk)->R_inv; DGEMM(A, R_inv)", + "GEQP3(R_sk)=Q*R_buf*P^T; R_inv=P*TRSM(R_buf,Q^T); DGEMM(A, R_inv)", + "GESDD(R_sk)=U*S*Vt; R_inv=V*diag(1/S)*U^T; DGEMM(A, R_inv)", + "BQRRP(R_sk)=Q*R_buf*P^T; R_inv=P*TRSM(R_buf,Q^T); DGEMM(A, R_inv)", +}; + +// ============================================================================ +// Helpers +// ============================================================================ + +template +static T rel_diff(const T* A, const T* B, int64_t len) { + T nd = 0, na = 0; + for (int64_t i = 0; i < len; ++i) { + T d = A[i] - B[i]; + nd += d * d; na += A[i] * A[i]; + } + return (na > 0) ? std::sqrt(nd / na) : std::sqrt(nd); +} + +// Condition number of the Gram matrix G = A_pre^T A_pre. Uses syrk +// (upper triangle) + symmetrize so gesdd inside cond_num_check sees a +// full symmetric matrix. +template +static T gram_condition_number(const T* A_pre, int64_t m, int64_t n) { + std::vector G(n * n, 0.0); + blas::syrk(Layout::ColMajor, Uplo::Upper, Op::Trans, + n, m, (T)1.0, A_pre, m, (T)0.0, G.data(), n); + RandBLAS::symmetrize(Layout::ColMajor, Uplo::Upper, n, G.data(), n); + return RandLAPACK::util::cond_num_check(n, n, G.data(), /*verbose=*/false); +} + +// Full CQRRT pipeline (matching CQRRT_linops Gram computation): +// G = A_orig^T * A_pre (GEMM, full n×n, not exploiting symmetry) +// G = (R_sketch)^{-T} * G (TRSM Left, backward-stable left factor on original R^sk) +// Zero lower triangle of G (POTRF/TRMM only use upper triangle; lower has TRSM output) +// R_chol = chol(G) (POTRF, upper triangle only) +// R_final = R_chol * R_sketch (TRMM) +// Q = A_orig * R_final^{-1} (TRSM on copy of A_orig) +// return orth_error(Q) +// Does NOT modify A_pre or A_orig. +template +static T cholqr_orth_error(const std::vector& A_pre, const T* A_orig, + int64_t m, int64_t n, const T* R_sketch) { + std::vector G(n * n, 0.0); + blas::gemm(Layout::ColMajor, Op::Trans, Op::NoTrans, + n, n, m, (T)1.0, A_orig, m, A_pre.data(), m, (T)0.0, G.data(), n); + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::Trans, Diag::NonUnit, + n, n, (T)1.0, R_sketch, n, G.data(), n); + // Zero strictly lower triangle: TRSM fills the full n×n matrix; the subsequent + // TRMM(Right,Upper) reads lower-triangle entries of G when computing upper-triangle + // output entries and would produce corrupted R_chol if left non-zero. + if (n > 1) + lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, &G.data()[1], n); + if (lapack::potrf(Uplo::Upper, n, G.data(), n)) + return std::numeric_limits::infinity(); + blas::trmm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, n, n, (T)1.0, R_sketch, n, G.data(), n); + std::vector Q(A_orig, A_orig + m * n); + blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, m, n, (T)1.0, G.data(), n, Q.data(), m); + return RandLAPACK::testing::orthogonality_error(Q.data(), m, n); +} + +// ============================================================================ +// One trial: N_PATHS-path orth comparison (shared sketch) + +// independent path [1] vs [2] step-by-step divergence +// ============================================================================ + +template +struct TrialResult { + // Per-path metrics (shared sketch, all N_PATHS paths) + double cond_Apre[N_PATHS]; + double cond_G[N_PATHS]; + double orth_Q[N_PATHS]; + // Cross-path relative differences of A_pre (reference = path [1], shared + // sketch): entry p holds rel_diff(Apre[0], Apre[p]); entry 0 is unused. + double rd_Apre_vs1[N_PATHS]; + // Step-by-step pipeline divergence: paths [1] vs [2], faithful to actual implementations + // Path [1] (CQRRT_expl): sketch via sketch_general(S, A_dense); TRSM in-place; SYRK + // Path [2] (CQRRT_linop): sketch via A_linop(Side::Right, S) [SpGEMM]; + // TRSM_IDENTITY → R_inv; A_linop fwd/adj; + // TRSM(Left,Trans) Gram completion on R_sk; TRMM R_final + // -1 = not computed: the potrf that would have produced this quantity failed. + double rd_Msk_12; // M^sk: raw sketch Ahat = S*A (different code paths) + double rd_Rsk_12; // R^sk: QR factor of the above + double rd_Apre_12_step; // MR^pre: TRSM in-place vs A_linop(fwd, R_inv) + double rd_G_12; // Gram: SYRK(A_pre) vs A_linop(adj, A_pre)+TRSM(R_sk) completion + double rd_Rchol_12; // R^chol: Cholesky factor + double rd_Rfinal_12; // R: R_final = R_chol * R_sk + // Sketch diagnostic + double cond_Rsk; +}; + +template +static TrialResult run_trial( + LinOpT& A_linop, + const T* A_dense, + int64_t m, int64_t n, + T d_factor, int64_t sketch_nnz, + RandBLAS::RNGState& state) +{ + TrialResult res{}; + // Part B early-returns on a Cholesky failure before every rd_* field is + // computed; a value-initialized 0.0 would then read as "paths agree + // exactly," exactly in the ill-conditioned regime this tool exists to + // probe. Seed the suite-wide -1 sentinel ("not computed, factorization + // failed") and let each Part B step overwrite it on success. + res.rd_Msk_12 = res.rd_Rsk_12 = res.rd_Apre_12_step = -1.0; + res.rd_G_12 = res.rd_Rchol_12 = res.rd_Rfinal_12 = -1.0; + int64_t d = (int64_t)std::ceil(d_factor * n); + + // Save RNG state before any sketching; Part B (step-by-step) uses this + // to compute independent sketches for paths [1] and [2]. + auto initial_state = state; + + // ---------------------------------------------------------------- + // Part A: Shared sketch → R_sk, N_PATHS-path orth comparison + // ---------------------------------------------------------------- + RandBLAS::SparseDist Ds(d, m, sketch_nnz, RandBLAS::Axis::Short); + RandBLAS::SparseSkOp S(Ds, state); + // Advance state past what S consumed: the BQRRP path below (Method E) + // also draws from `state` and must not reuse this sketch's seed material. + state = S.next_state; + std::vector Ahat(d * n, 0.0); + RandBLAS::sketch_general(Layout::ColMajor, Op::NoTrans, Op::NoTrans, + d, n, m, (T)1.0, S, A_dense, m, + (T)0.0, Ahat.data(), d); + + std::vector R_sk(n * n, 0.0); + { + std::vector tau(n); + lapack::geqrf(d, n, Ahat.data(), d, tau.data()); + for (int64_t j = 0; j < n; ++j) + for (int64_t i = 0; i <= j; ++i) + R_sk[i + j*n] = Ahat[i + j*d]; + } + res.cond_Rsk = (double)RandLAPACK::util::cond_num_check(n, n, R_sk.data(), /*verbose=*/false); + + // ---------------------------------------------------------------- + // Explicit inverses of R_sk via two methods + // ---------------------------------------------------------------- + + // Method A1: TRSM on identity, column-ordered solve (path [2]) + // Solve R_sk * X = I (Side::Left): each column of X is an independent + // backward-stable solve R_sk * x_j = e_j. This is the RandLAPACK default + // (PCholQRPrecondMethod::TRSM_IDENTITY, comps/rl_cholqr.hh), including + // the trailing lower-triangle laset the primitive performs. + std::vector R_inv_trsm_left(n * n, T(0)); + RandLAPACK::util::eye(n, n, R_inv_trsm_left.data()); + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, n, n, (T)1.0, + R_sk.data(), n, R_inv_trsm_left.data(), n); + if (n > 1) + lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, R_inv_trsm_left.data() + 1, n); + + // Method A2: TRSM on identity, reversed row-ordered solve (path [3]) + // Solve X * R_sk = I (Side::Right): rows of X carry independent + // perturbations of R_sk, and the error of the composite M * X scales + // with kappa(R_sk). Kept only to document the solve-ordering effect. + std::vector R_inv_trsm_right(n * n, T(0)); + RandLAPACK::util::eye(n, n, R_inv_trsm_right.data()); + blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, n, n, (T)1.0, + R_sk.data(), n, R_inv_trsm_right.data(), n); + if (n > 1) + lapack::laset(MatrixType::Lower, n-1, n-1, (T)0.0, (T)0.0, R_inv_trsm_right.data() + 1, n); + + // Method B: LAPACK trtri (path [4]) + std::vector R_inv_trtri(R_sk.begin(), R_sk.end()); + lapack::trtri(Uplo::Upper, Diag::NonUnit, n, R_inv_trtri.data(), n); + + // Method C: GEQP3 factorization of R_sk (path [5]) + // R_sk * P = Q_buf * R_buf (GEQP3) + // R_sk^{-1} = P * R_buf^{-1} * Q_buf^T + // Computed via Option A: ungqr (explicit Q) + TRSM (cheaper than trtri + ormqr) + // 1. ungqr -> Q_buf explicit (~4n^3/3 flops) + // 2. W = Q_buf^T (explicit transpose, O(n^2)) + // 3. TRSM(Left, R_buf, W) -> W = R_buf^{-1} * Q_buf^T (~n^3/2 flops) + // 4. scatter W by jpiv -> R_sk^{-1} = P * W + // Total: ~11n^3/6. Alternative (trtri + ormqr) costs ~7n^3/3. + std::vector R_inv_geqp3(n * n, 0.0); + { + std::vector R_copy(R_sk.begin(), R_sk.end()); + std::vector jpiv(n, 0); + std::vector tau_qr(n); + lapack::geqp3(n, n, R_copy.data(), n, jpiv.data(), tau_qr.data()); + + // Extract upper triangular R_buf before overwriting with Q + std::vector R_buf(n * n, 0.0); + lapack::lacpy(MatrixType::Upper, n, n, R_copy.data(), n, R_buf.data(), n); + + // Expand Q_buf from Householder reflectors (overwrites R_copy) + lapack::ungqr(n, n, n, R_copy.data(), n, tau_qr.data()); + + // W = R_buf^{-1} * Q_buf^T via TRSM: initialize W as Q_buf^T, then solve in-place + std::vector W(n * n, 0.0); + for (int64_t i = 0; i < n; ++i) + for (int64_t j = 0; j < n; ++j) + W[i + j*n] = R_copy[j + i*n]; // W := Q_buf^T (col-major) + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, n, n, (T)1.0, + R_buf.data(), n, W.data(), n); + + // R_sk^{-1} = P * W: row (jpiv[k]-1) of R_inv gets row k of W + for (int64_t k = 0; k < n; ++k) + for (int64_t j = 0; j < n; ++j) + R_inv_geqp3[(jpiv[k]-1) + j*n] = W[k + j*n]; + } + + // Method D: SVD of R_sk (path [6]) + // R_sk = U * diag(s) * Vt + // R_sk^{-1} = V * diag(1/s) * U^T = Vt^T * diag(1/s) * U^T + std::vector R_inv_svd(n * n, 0.0); + { + std::vector R_copy(R_sk.begin(), R_sk.end()); + std::vector U(n * n, 0.0), Vt(n * n, 0.0), s(n); + lapack::gesdd(lapack::Job::AllVec, n, n, R_copy.data(), n, + s.data(), U.data(), n, Vt.data(), n); + // Scale row k of Vt by 1/s[k]: Vt[k + j*n] is row k, col j (col-major) + for (int64_t k = 0; k < n; ++k) + for (int64_t j = 0; j < n; ++j) + Vt[k + j*n] /= s[k]; + // R_inv = scaled_Vt^T * U^T + blas::gemm(Layout::ColMajor, Op::Trans, Op::Trans, + n, n, n, (T)1.0, Vt.data(), n, U.data(), n, + (T)0.0, R_inv_svd.data(), n); + } + + // Method E: BQRRP factorization of R_sk (path [7]) + // Same output format as GEQP3: R_sk * P = Q_buf * R_buf, then + // R_sk^{-1} = P * R_buf^{-1} * Q_buf^T. + // BQRRP uses blocked randomized QRCP; block size matches CQRRT_linops + // adaptive heuristic (1.0 for n <= 2000, 0.5 for n <= 8000, 1/32 else). + std::vector R_inv_bqrrp(n * n, 0.0); + { + std::vector R_copy(R_sk.begin(), R_sk.end()); + std::vector jpiv(n, 0); + std::vector tau_qr(n); + + T block_ratio; + if (n <= 2000) block_ratio = (T)1.0; + else if (n <= 8000) block_ratio = (T)0.5; + else block_ratio = (T)1.0 / (T)32; + int64_t bqrrp_block = std::max(1, (int64_t)(n * block_ratio)); + RandLAPACK::BQRRP bqrrp(false, bqrrp_block); + bqrrp.call(n, n, R_copy.data(), n, (T)1.0, tau_qr.data(), jpiv.data(), state); + + std::vector R_buf(n * n, 0.0); + lapack::lacpy(MatrixType::Upper, n, n, R_copy.data(), n, R_buf.data(), n); + + lapack::ungqr(n, n, n, R_copy.data(), n, tau_qr.data()); + + std::vector W(n * n, 0.0); + for (int64_t i = 0; i < n; ++i) + for (int64_t j = 0; j < n; ++j) + W[i + j*n] = R_copy[j + i*n]; // W := Q_buf^T (col-major) + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, n, n, (T)1.0, + R_buf.data(), n, W.data(), n); + + for (int64_t k = 0; k < n; ++k) + for (int64_t j = 0; j < n; ++j) + R_inv_bqrrp[(jpiv[k]-1) + j*n] = W[k + j*n]; + } + + // ---------------------------------------------------------------- + // Compute all N_PATHS preconditioned matrices: + // Apre[0]: TRSM in-place (path [1], CQRRT_expl) + // Apre[p]: GEMM + R_invs[p-1], p >= 1 (paths [2]..[7]) + // ---------------------------------------------------------------- + std::array, N_PATHS> Apre; + for (auto& a : Apre) a.resize(m * n, T(0)); + + Apre[0].assign(A_dense, A_dense + m*n); + blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, m, n, (T)1.0, R_sk.data(), n, Apre[0].data(), m); + + const T* R_invs[N_PATHS - 1] = { + R_inv_trsm_left.data(), R_inv_trsm_right.data(), R_inv_trtri.data(), + R_inv_geqp3.data(), R_inv_svd.data(), R_inv_bqrrp.data() + }; + for (int p = 1; p < N_PATHS; ++p) + blas::gemm(Layout::ColMajor, Op::NoTrans, Op::NoTrans, + m, n, n, (T)1.0, A_dense, m, R_invs[p-1], n, (T)0.0, Apre[p].data(), m); + + // ---------------------------------------------------------------- + // Cross-path relative differences (reference = path [1] = Apre[0]) + // ---------------------------------------------------------------- + for (int p = 1; p < N_PATHS; ++p) + res.rd_Apre_vs1[p] = (double)rel_diff(Apre[0].data(), Apre[p].data(), m*n); + + // ---------------------------------------------------------------- + // Part B: Independent step-by-step divergence, paths [1] vs [2] + // + // Both paths compute their own sketch and R_sk from initial_state + // (same seed → same result, but as separate objects). + // + // Path [1] (CQRRT_expl): TRSM in-place on A; Gram via SYRK. + // Path [2] (CQRRT_linop, block_size=0): + // R_inv = TRSM_IDENTITY(R_sk); + // A_pre = GEMM(A, R_inv) [fwd linop call] + // G = GEMM(A^T, A_pre) [adj linop call] + // G = TRSM(Left,Trans,R_sk,G) [complete Gram: (R_sk)^{-T} * A^T * A * R_inv, + // backward-stable solve on the original R_sk] + // ---------------------------------------------------------------- + + // Run Part B as a lambda so early returns on Cholesky failure are clean. + [&]() { + // ---- Step 1: Sketch ---- + // Path [1] (CQRRT_expl): sketch_general(S, A_dense), left SPMM on dense copy + RandBLAS::SparseDist Ds_1(d, m, sketch_nnz, RandBLAS::Axis::Short); + RandBLAS::SparseSkOp S_1(Ds_1, initial_state); + std::vector Ahat_1(d * n, 0.0); + RandBLAS::sketch_general(Layout::ColMajor, Op::NoTrans, Op::NoTrans, + d, n, m, (T)1.0, S_1, A_dense, m, + (T)0.0, Ahat_1.data(), d); + + // Path [2] (CQRRT_linop): A_linop(Side::Right, S), SpGEMM on sparse CSR matrix + RandBLAS::SparseDist Ds_2(d, m, sketch_nnz, RandBLAS::Axis::Short); + RandBLAS::SparseSkOp S_2(Ds_2, initial_state); // same seed → same S + std::vector Ahat_2(d * n, 0.0); + A_linop(Side::Right, Layout::ColMajor, Op::NoTrans, Op::NoTrans, + d, n, m, (T)1.0, S_2, (T)0.0, Ahat_2.data(), d); + + // M^sk diff: before geqrf overwrites the sketch buffers + res.rd_Msk_12 = (double)rel_diff(Ahat_1.data(), Ahat_2.data(), d*n); + + // ---- Step 2: QR → R_sk ---- + std::vector R_sk_1(n * n, 0.0); + { + std::vector tau_1(n); + lapack::geqrf(d, n, Ahat_1.data(), d, tau_1.data()); + for (int64_t j = 0; j < n; ++j) + for (int64_t i = 0; i <= j; ++i) + R_sk_1[i + j*n] = Ahat_1[i + j*d]; + } + std::vector R_sk_2(n * n, 0.0); + { + std::vector tau_2(n); + lapack::geqrf(d, n, Ahat_2.data(), d, tau_2.data()); + for (int64_t j = 0; j < n; ++j) + for (int64_t i = 0; i <= j; ++i) + R_sk_2[i + j*n] = Ahat_2[i + j*d]; + } + res.rd_Rsk_12 = (double)rel_diff(R_sk_1.data(), R_sk_2.data(), n*n); + + // ---- Path [1]: CQRRT_expl, TRSM in-place, SYRK ---- + std::vector Apre_1(A_dense, A_dense + m*n); + blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, m, n, (T)1.0, R_sk_1.data(), n, Apre_1.data(), m); + + std::vector G_1(n*n, 0.0); + blas::syrk(Layout::ColMajor, Uplo::Upper, Op::Trans, + n, m, (T)1.0, Apre_1.data(), m, (T)0.0, G_1.data(), n); + for (int64_t j = 0; j < n; ++j) + for (int64_t i = j+1; i < n; ++i) + G_1[i + j*n] = G_1[j + i*n]; + + std::vector Rchol_1(G_1); + if (lapack::potrf(Uplo::Upper, n, Rchol_1.data(), n)) return; + for (int64_t j = 0; j < n; ++j) + for (int64_t i = j+1; i < n; ++i) + Rchol_1[i + j*n] = (T)0.0; + + std::vector Rfinal_1(Rchol_1); + blas::trmm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, n, n, (T)1.0, R_sk_1.data(), n, Rfinal_1.data(), n); + + // ---- Path [2]: CQRRT_linop, TRSM_IDENTITY, linop fwd/adj, TRSM Gram completion ---- + // R_inv via TRSM_IDENTITY: solve R_sk_2 * X = I (Side::Left), matching + // cholqr_primitive's shipping default (comps/rl_cholqr.hh). + std::vector R_inv_2(n * n, T(0)); + RandLAPACK::util::eye(n, n, R_inv_2.data()); + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, n, n, (T)1.0, R_sk_2.data(), n, R_inv_2.data(), n); + for (int64_t j = 0; j < n; ++j) + for (int64_t i = j+1; i < n; ++i) + R_inv_2[i + j*n] = (T)0.0; + + // fwd: Apre_2 = A * R_inv_2 via A_linop(Side::Left, NoTrans) + std::vector Apre_2(m * n, 0.0); + A_linop(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, + m, n, n, (T)1.0, R_inv_2.data(), n, (T)0.0, Apre_2.data(), m); + res.rd_Apre_12_step = (double)rel_diff(Apre_1.data(), Apre_2.data(), m*n); + + // adj: G_2 = A^T * Apre_2 via A_linop(Side::Left, Trans) + std::vector G_2(n * n, 0.0); + A_linop(Side::Left, Layout::ColMajor, Op::Trans, Op::NoTrans, + n, n, m, (T)1.0, Apre_2.data(), m, (T)0.0, G_2.data(), n); + // Complete Gram: G_2 = (R_sk_2)^{-T} * G_2 (backward-stable TRSM on original R_sk) + blas::trsm(Layout::ColMajor, Side::Left, Uplo::Upper, Op::Trans, Diag::NonUnit, + n, n, (T)1.0, R_sk_2.data(), n, G_2.data(), n); + res.rd_G_12 = (double)rel_diff(G_1.data(), G_2.data(), n*n); + + std::vector Rchol_2(G_2); + if (lapack::potrf(Uplo::Upper, n, Rchol_2.data(), n)) return; + for (int64_t j = 0; j < n; ++j) + for (int64_t i = j+1; i < n; ++i) + Rchol_2[i + j*n] = (T)0.0; + res.rd_Rchol_12 = (double)rel_diff(Rchol_1.data(), Rchol_2.data(), n*n); + + std::vector Rfinal_2(Rchol_2); + blas::trmm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, + Diag::NonUnit, n, n, (T)1.0, R_sk_2.data(), n, Rfinal_2.data(), n); + res.rd_Rfinal_12 = (double)rel_diff(Rfinal_1.data(), Rfinal_2.data(), n*n); + }(); + + // ---------------------------------------------------------------- + // Per-path metrics + // ---------------------------------------------------------------- + for (int p = 0; p < N_PATHS; ++p) { + res.cond_Apre[p] = (double)RandLAPACK::util::cond_num_check(m, n, Apre[p].data(), /*verbose=*/false); + res.cond_G[p] = (double)gram_condition_number(Apre[p].data(), m, n); + res.orth_Q[p] = (double)cholqr_orth_error(Apre[p], A_dense, m, n, R_sk.data()); + } + + return res; +} + +// ============================================================================ +// Shared: write CSV header and run trials given a dense matrix +// ============================================================================ + +template +static void write_csv_and_run( + LinOpT& A_linop, + const std::vector& A_dense, + int64_t m, int64_t n, + T d_factor, int64_t sketch_nnz, int64_t num_runs, + T cond_A, double kappa_target, // kappa_target < 0 means "from file" + const std::string& matrix_label, + const std::string& output_dir, + const std::string& argv_line) +{ + std::string ts = make_run_timestamp(); + std::string csv_path = output_dir + "/diagnostic_" + ts + ".csv"; + std::ofstream csv(csv_path); + + csv << "# CQRRT Preconditioner Comparison\n"; + csv << "# Date: " << ts << "\n"; + csv << "# host: " << get_hostname() << "\n"; + csv << "# argv: " << argv_line << "\n"; + csv << "# RANDLAPACK_GIT_COMMIT=" << RandLAPACK::bench::env_or("RANDLAPACK_GIT_COMMIT") << "\n"; + csv << "# This tool hand-rolls its own CQRRT pipeline rather than calling\n"; + csv << "# cholqr_primitive, so the RANDLAPACK_GRAM_LEFT / RANDLAPACK_SCHOLQR3_SHIFT /\n"; + csv << "# RANDLAPACK_BLAS2_THREADS / RANDLAPACK_FFT_THREADS / RANDLAPACK_SOLVE_FFT_MATCH\n"; + csv << "# knobs that steer the shipped drivers do not affect any column below.\n"; + csv << "# Matrix: " << matrix_label << "\n"; + csv << "# m=" << m << " n=" << n << " d_factor=" << d_factor + << " sketch_nnz=" << sketch_nnz << "\n"; + csv << "# cond_A=" << std::scientific << std::setprecision(6) << cond_A; + if (kappa_target > 0) + csv << " kappa_target=" << std::scientific << std::setprecision(6) << kappa_target; + csv << "\n"; + for (int p = 0; p < N_PATHS; ++p) + csv << "# path " << (p+1) << ": " << PATH_NAMES[p] << "\n"; + csv << "# rd_Msk_12/rd_Rsk_12/rd_Apre_12_step/rd_G_12/rd_Rchol_12/rd_Rfinal_12:\n"; + csv << "# -1 = not computed; Part B's Cholesky factorization failed before this\n"; + csv << "# quantity was formed (no adaptive-shift retry is modeled here, so this is\n"; + csv << "# expected in the high-kappa regime this tool probes).\n"; + csv << "run,"; + for (int p = 1; p <= N_PATHS; ++p) csv << "orth_Q" << p << ","; + for (int p = 1; p <= N_PATHS; ++p) csv << "cond_Apre" << p << ","; + for (int p = 1; p <= N_PATHS; ++p) csv << "cond_G" << p << ","; + for (int p = 2; p <= N_PATHS; ++p) csv << "rd_Apre_1" << p << ","; + csv << "rd_Msk_12,rd_Rsk_12,rd_Apre_12_step,rd_G_12,rd_Rchol_12,rd_Rfinal_12," + << "cond_Rsk\n"; + + RandBLAS::RNGState base_state(42); + for (int64_t r = 0; r < num_runs; ++r) { + auto state = base_state; + if (r > 0) state.key.incr(r); + + auto res = run_trial(A_linop, A_dense.data(), m, n, d_factor, sketch_nnz, state); + + printf(" run %lld orth_error(Q = A * R_final^{-1}):\n", (long long)r); + for (int p = 0; p < N_PATHS; ++p) + printf(" [%d] %-18s %12.3e\n", p+1, PATH_NAMES[p], res.orth_Q[p]); + + printf(" run %lld cond(MR^pre):\n", (long long)r); + for (int p = 0; p < N_PATHS; ++p) + printf(" [%d] %-18s %12.3e\n", p+1, PATH_NAMES[p], res.cond_Apre[p]); + + printf(" run %lld cond(G = MR^pre^T MR^pre):\n", (long long)r); + for (int p = 0; p < N_PATHS; ++p) + printf(" [%d] %-18s %12.3e\n", p+1, PATH_NAMES[p], res.cond_G[p]); + + printf(" run %lld rel_diff(MR^pre) vs [1]:\n", (long long)r); + for (int p = 1; p < N_PATHS; ++p) + printf(" rd_1%d (%-19s): %12.3e\n", p+1, PATH_NAMES[p], res.rd_Apre_vs1[p]); + + printf(" run %lld step-by-step divergence [1] vs [2] (expl: sketch_general; linop: SpGEMM):\n", (long long)r); + printf(" M^sk: %12.3e\n", res.rd_Msk_12); + printf(" R^sk: %12.3e\n", res.rd_Rsk_12); + printf(" MR^pre: %12.3e\n", res.rd_Apre_12_step); + printf(" G: %12.3e\n", res.rd_G_12); + printf(" R^chol: %12.3e\n", res.rd_Rchol_12); + printf(" R: %12.3e\n", res.rd_Rfinal_12); + + printf(" run %lld cond(R_sk): %9.3e\n\n", (long long)r, res.cond_Rsk); + + csv << r << "," << std::scientific << std::setprecision(6); + for (int p = 0; p < N_PATHS; ++p) csv << res.orth_Q[p] << ","; + for (int p = 0; p < N_PATHS; ++p) csv << res.cond_Apre[p] << ","; + for (int p = 0; p < N_PATHS; ++p) csv << res.cond_G[p] << ","; + for (int p = 1; p < N_PATHS; ++p) csv << res.rd_Apre_vs1[p] << ","; + csv << res.rd_Msk_12 << "," << res.rd_Rsk_12 << "," << res.rd_Apre_12_step << "," + << res.rd_G_12 << "," << res.rd_Rchol_12 << "," << res.rd_Rfinal_12 << "," + << res.cond_Rsk << "\n"; + } + csv.close(); + + std::cout << " Legend:\n"; + for (int p = 0; p < N_PATHS; ++p) + printf(" [%d] %-18s %s\n", p+1, PATH_NAMES[p], PATH_DESCS[p]); + std::cout << "\n CSV written to: " << csv_path << "\n"; +} + +// ============================================================================ +// Main benchmark +// ============================================================================ + +template +int run_benchmark(int argc, char* argv[]) { + // argc < 3 (not < 2): output_dir is read from argv[2] unconditionally + // right below, and argc==2 (e.g. just ) would read past argv's end. + if (argc < 3) { + std::cerr << "Usage (file mode): " << argv[0] + << " [sketch_nnz]\n" + << "Usage (generate mode): " << argv[0] + << " gen [sketch_nnz]\n"; + return 1; + } + + std::string output_dir = argv[2]; + std::string argv_line = quote_join_argv(argc, argv); + bool is_generate = (argc >= 4 && std::string(argv[3]) == "gen"); + + if (is_generate) { + // generate mode: prec output_dir gen m n kappa density d_factor runs [sketch_nnz] + // Required form is argv[0..9] (10 tokens); argc < 10 rejects it. The + // trailing [sketch_nnz] is argv[10] and requires argc >= 11. + if (argc < 10) { + std::cerr << "Usage (generate mode): " << argv[0] + << " gen [sketch_nnz]\n"; + return 1; + } + int64_t m = std::stoll(argv[4]); + int64_t n = std::stoll(argv[5]); + T kappa = (T)std::stod(argv[6]); + T density = (T)std::stod(argv[7]); + T d_factor = (T)std::stod(argv[8]); + int64_t num_runs = std::stoll(argv[9]); + int64_t sketch_nnz = (argc >= 11) ? std::stoll(argv[10]) : 4; + + if (num_runs < 1) { + std::cerr << "Error: runs must be >= 1 (got " << num_runs << ")\n"; + return 1; + } + + std::cout << "\n=== CQRRT Preconditioner Comparison (generate mode) ===\n"; + std::cout << " Size: " << m << " x " << n << "\n"; + std::cout << " kappa: " << std::scientific << std::setprecision(3) << (double)kappa << "\n"; + std::cout << " density: " << density << "\n"; + std::cout << " d_factor: " << d_factor << "\n"; + std::cout << " sketch_nnz: " << sketch_nnz << "\n"; + std::cout << " runs: " << num_runs << "\n"; +#ifdef _OPENMP + std::cout << " OMP threads: " << omp_get_max_threads() << "\n"; +#endif + + RandBLAS::RNGState gen_state(0); + auto A_coo = RandLAPACK::gen::gen_sparse_cond_coo(m, n, kappa, gen_state, density); + RandBLAS::sparse_data::csr::CSRMatrix A_csr(m, n); + RandBLAS::sparse_data::conversions::coo_to_csr(A_coo, A_csr); + RandLAPACK::linops::SparseLinOp> A_linop(m, n, A_csr); + + std::vector A_dense(m * n, 0.0); + { + std::vector Eye(n * n, T(0)); + RandLAPACK::util::eye(n, n, Eye.data()); + A_linop(Layout::ColMajor, Op::NoTrans, Op::NoTrans, + m, n, n, (T)1.0, Eye.data(), n, (T)0.0, A_dense.data(), m); + } + T cond_A = RandLAPACK::util::cond_num_check(m, n, A_dense.data(), /*verbose=*/false); + std::cout << " cond(A): " << std::scientific << std::setprecision(3) << (double)cond_A << "\n\n"; + + std::string label = "gen_" + std::to_string(m) + "x" + std::to_string(n) + + "_kappa" + std::to_string((int)std::round(std::log10((double)kappa))); + write_csv_and_run(A_linop, A_dense, m, n, d_factor, sketch_nnz, num_runs, + cond_A, (double)kappa, label, output_dir, argv_line); + } else { + // file mode: prec output_dir mtx_path d_factor runs [sketch_nnz] + if (argc < 6) { + std::cerr << "Usage (file mode): " << argv[0] + << " [sketch_nnz]\n"; + return 1; + } + std::string mtx_path = argv[3]; + T d_factor = (T)std::stod(argv[4]); + int64_t num_runs = std::stoll(argv[5]); + int64_t sketch_nnz = (argc >= 7) ? std::stoll(argv[6]) : 4; + + if (num_runs < 1) { + std::cerr << "Error: runs must be >= 1 (got " << num_runs << ")\n"; + return 1; + } + + int64_t m, n, nnz; + auto csr = load_csr(mtx_path, m, n, nnz); + RandLAPACK::linops::SparseLinOp> A_linop(m, n, csr); + + std::vector A_dense(m * n, 0.0); + { + std::vector Eye(n * n, T(0)); + RandLAPACK::util::eye(n, n, Eye.data()); + A_linop(Layout::ColMajor, Op::NoTrans, Op::NoTrans, + m, n, n, (T)1.0, Eye.data(), n, (T)0.0, A_dense.data(), m); + } + T cond_A = RandLAPACK::util::cond_num_check(m, n, A_dense.data(), /*verbose=*/false); + int64_t d = (int64_t)std::ceil(d_factor * n); + + std::cout << "\n=== CQRRT Preconditioner Comparison ===\n"; + std::cout << " Matrix: " << mtx_path << "\n"; + std::cout << " Size: " << m << " x " << n << " (nnz=" << nnz << ")\n"; + std::cout << " d_factor: " << d_factor << " (d=" << d << ")\n"; + std::cout << " sketch_nnz: " << sketch_nnz << "\n"; + std::cout << " runs: " << num_runs << "\n"; + std::cout << " cond(A): " << std::scientific << std::setprecision(3) << (double)cond_A << "\n"; +#ifdef _OPENMP + std::cout << " OMP threads: " << omp_get_max_threads() << "\n"; +#endif + std::cout << "\n"; + + write_csv_and_run(A_linop, A_dense, m, n, d_factor, sketch_nnz, num_runs, + cond_A, -1.0, mtx_path, output_dir, argv_line); + } + + return 0; +} + +int main(int argc, char* argv[]) { + if (argc < 2) { + std::cerr << "Usage (file mode): " << argv[0] + << " [sketch_nnz]\n" + << "Usage (generate mode): " << argv[0] + << " gen [sketch_nnz]\n"; + return 1; + } + std::string prec = argv[1]; + if (prec == "double") return run_benchmark(argc, argv); + if (prec == "float") return run_benchmark(argc, argv); + std::cerr << "Unknown precision '" << prec << "' (use double or float)\n"; + return 1; +} diff --git a/benchmark/bench_CQRRT_linops/CQRRT_linop_applications.cc b/benchmark/bench_CQRRT_linops/CQRRT_linop_applications.cc new file mode 100644 index 000000000..bb755a9e3 --- /dev/null +++ b/benchmark/bench_CQRRT_linops/CQRRT_linop_applications.cc @@ -0,0 +1,2413 @@ +// Unified Q-less QR benchmark: IR-LSQ application, plus rspec (Algorithm 4). +// +// Pipeline: +// 1. Load matrices (FEM mode: K, M, V .mtx files; sparse mode: a single A.mtx). +// 2. (FEM only) Cholesky-factorize M = L L^T via CholSolverLinOp(half_solve=true). +// 3. (FEM only) Build J = L^{-1} K V as a doubly-nested CompositeOperator +// J = CompositeOperator(L_inv_op, CompositeOperator(K_op, V_op)). +// 4. Run Q-less QR via one of 5 variants (CQRRT_linop, CholQR, sCholQR3, +// sCholQR3_basic, CholQR2), selected by method_mask. +// 5. Post-processing dictated by : +// irlsq: IterRefineLSQ from x_0 = 0 (no sketch-and-solve initial guess; +// the only sketch is S_1 inside Q-less QR, which produces R) +// rspec: reduced spectral approximation (Algorithm 4): Rayleigh-Ritz on +// range(C^j V_FEM), C = L^T (K - ω M)^{-1} L. FEM-only. +// +// Usage: +// ./CQRRT_linop_applications +// sparse [nnz] [b] [compute_cond] [method_mask] [noise_level] +// ./CQRRT_linop_applications +// [nnz] [b] [compute_cond] [method_mask] [noise_level] [omega] [power_j] +// +// mode = "irlsq" | "irlsq_reg" | "rspec" +// (main() hard-errors on an unrecognized mode; matching nothing, writing no +// CSV, and exiting 0 is a silent failure mode that has burned SLURM scripts.) +// method_mask = bitmask of methods (default 0b11111 = 31) +// bit 0 ( 1): CQRRT_linop (TRSM_IDENTITY) +// bit 1 ( 2): CholQR +// bit 2 ( 4): sCholQR3 +// bit 3 ( 8): sCholQR3_basic +// bit 4 ( 16): CholQR2 +// bit 5 ( 32): Blendenpik, published (warm + cold rows; not in the +// default mask) +// bit 6 ( 64): Blendenpik refined by the shared engine (warm + cold +// rows; see benchmark/refined_blendenpik.hh). NOTE this +// bit was previously CQRRT_linop_bqrrp and was reused +// for refine rows, so a script written against the old +// assignment gets refine rows, not the BQRRP variant. +// rspec mode accepts bits 0-4 only and warns on 32/64. +// The campaign mask 127 = bits 0-6 (all five Q-less methods + both Blendenpik +// families). +// +// Trailing optional args after precond_prec (irlsq / irlsq_reg): +// [ir_max_inner] inner-CG iteration cap per outer refinement step (default 200). +// With 2 outer steps this is what produced the fixed 400-iteration +// ceiling in earlier CSVs. Pass <= 0 to keep the default. +// [ir_inner_tol] inner-CG relative-residual tolerance (default: eps^0.85 in the +// working precision, ~4.9e-14 in double). Pass < 0 to keep it. +// [ir_round_drop] per-round inner-CG residual drop (default 1e-4; restart +// pacing, replacing [ir_inner_restarts] in this slot). +// Each round's CG stops after this relative drop and the outer +// loop restarts against the TRUE residual; ir_inner_tol survives +// as the absolute floor at which rounds stop immediately. Pass 0 +// for legacy fixed-tolerance rounds. Values >= 1 are rejected so +// stale scripts passing the old restart count fail loudly. +// [ir_n_steps] outer-round cap (default 20; previously 4). Under the +// paced scheme rounds are shallow and ir_outer_tol exits early, +// so strong preconditioners use a few rounds and weak ones get +// room to descend instead of being budget-truncated. +// [ir_outer_tol] outer early-exit tolerance on ||b - Jx||/||b|| (default < 0 => +// 10*eps of the solve precision; pass 0 to always run all steps). +// Makes the outer loop "refine until done, capped at ir_n_steps", +// the same contract as the Toeplitz benchmark's pcg_ne solver. +// These exist so a diagnostic sweep can separate "CG stagnates below an unreachable +// tolerance" from "CG is still converging when the cap stops it" without a rebuild. +// The per-run answer is written to the CSV as ir_inner_capped / ir_inner_relres / +// ir_inner_best_relres / ir_inner_best_iter. +// +// Warm-start policy: the sketch-and-solve x0 warm start is Blendenpik-only. +// Method mask bit 32 runs TWO variants, "Blendenpik" (its own warm start, the +// published configuration) and "Blendenpik_cold" (x0 = 0), and IterRefineLSQ +// always starts from x0 = 0 (per collaborator request). The former +// [ir_warm_start] and [bp_warm_start] CLI knobs are removed: warm x0 is +// Blendenpik's forward-error edge, not IterRefineLSQ's. + +#include "RandLAPACK.hh" +#include "rl_blaspp.hh" +#include "rl_lapackpp.hh" +#include "rl_gen.hh" + +#include +#include +#include +#include +#include +#include +#include +#ifdef _OPENMP +#include +#endif +#include +#include +#include +#include +#include + +// Extras utilities (Eigen-dependent) +#include "../../extras/misc/ext_util.hh" +#include "../../extras/misc/ext_sparse_axpy.hh" +#include "../../extras/linops/ext_cholsolver_linop.hh" +#include "RandLAPACK/testing/rl_test_utils.hh" +#include "cqrrt_bench_common.hh" + +// Linops algorithms +#include "rl_cholqr_linops.hh" +#include "rl_scholqr3_linops.hh" +#include "rl_blendenpik.hh" +#include "../refined_blendenpik.hh" +#include "RandLAPACK/testing/rl_memory_tracker.hh" + +using std::chrono::steady_clock; +using std::chrono::duration_cast; +using std::chrono::microseconds; + +// Helper families shared with bench_toeplitz_ls (cqrrt_bench_common.hh): +// stop-reason maps, env provenance, the rounds-CSV schema/row writer, the +// chol-shift fold, and the power-2-norm and blocked-orth-error estimators. +using RandLAPACK::bench::pcg_stop_reason; +using RandLAPACK::bench::lsqr_stop_reason; +using RandLAPACK::bench::write_env_line; +using RandLAPACK::bench::kRoundsCsvHeader; +using RandLAPACK::bench::write_round_row; +using RandLAPACK::bench::fold_chol_shift; +using RandLAPACK::bench::bench_chol_max_retries; +using RandLAPACK::bench::estimate_op_2norm; +using RandLAPACK::bench::compute_orth_error_explicit; +using RandLAPACK::bench::quote_join_argv; +using RandLAPACK::bench::get_hostname; +using RandLAPACK::bench::write_host_line; + +// ============================================================================ +// Condition-number injection + precision-casting helpers (irlsq_reg mode) +// ============================================================================ + +// Geometric column-scaling diagonal d[j] = kappa^(j/(n-1)), j = 0..n-1, so the +// column-norm spread injected into J = L^{-1} K (V D) is kappa (d[0]=1, d[n-1]=kappa). +// kappa <= 1 (or n <= 1) means no scaling (all ones), i.e. native conditioning. +static std::vector geometric_colscale(int64_t n, double kappa) { + std::vector d(n, 1.0); + if (kappa > 1.0 && n > 1) { + for (int64_t j = 0; j < n; ++j) + d[j] = std::pow(kappa, (double)j / (double)(n - 1)); + } + return d; +} + +// Right-multiply a CSR matrix by diag(d): column j is scaled by d[j]. +// In CSR, nonzero p sits in column colidxs[p], so vals[p] *= d[colidxs[p]]. +template +static void scale_csr_columns(RandBLAS::sparse_data::csr::CSRMatrix& A, + const std::vector& d) { + for (int64_t p = 0; p < A.nnz; ++p) + A.vals[p] = (T)((double)A.vals[p] * d[(int64_t)A.colidxs[p]]); +} + +// Cast a CSR matrix to a different value precision (structure copied, values cast). +template +static RandBLAS::sparse_data::csr::CSRMatrix +csr_cast(const RandBLAS::sparse_data::csr::CSRMatrix& src) { + RandBLAS::sparse_data::csr::CSRMatrix dst(src.n_rows, src.n_cols); + if (src.nnz > 0) { + dst.reserve(src.nnz); + std::copy(src.rowptr, src.rowptr + src.n_rows + 1, dst.rowptr); + std::copy(src.colidxs, src.colidxs + src.nnz, dst.colidxs); + for (int64_t p = 0; p < src.nnz; ++p) dst.vals[p] = (Tdst)src.vals[p]; + } + return dst; +} + +// Unit roundoff u = eps/2 in precision P (collaborator's mu = mu_factor * u). +template +static P unit_roundoff() { return std::numeric_limits

::epsilon() / (P)2; } + +// ============================================================================ +// Result struct (unified: sentinel values for fields irrelevant to the mode) +// ============================================================================ + +template +struct bench_result { + int64_t m, n; + int64_t run_idx; + std::string alg_name; + T noise_level; + + long chol_time_us; // FEM: shared, measured once. Sparse: 0. + int qr_status; // 0 on success + long qr_time_us; // -1 if QR failed + // -1 = no Cholesky ran in this row (Blendenpik family, unpreconditioned, + // never overwrites this field) or QR failed before a retry count existed. + // 0 = Cholesky ran unshifted. The five Q-less branches always overwrite + // this on both success and failure, so real Cholesky rows are unaffected. + int chol_retries = -1; + // The shift the retries actually applied. -1 = no Cholesky in this row + // (Blendenpik family) or QR failed before a shift record existed; 0 = + // every pass unshifted. abs = pass-1 absolute shift, i.e. the + // Tikhonov-equivalent regularization baked into the returned R; rel = max + // over passes of shift/trace(G), the scale-free severity. A bare retry + // count cannot separate a rounding-level rescue from a spectrum-truncating + // one; these can. + T chol_shift_abs = -1; + T chol_shift_rel = -1; + + // Q-factor orthogonality: ||Q^T Q - I||_F / sqrt(n), computed for all methods. + T orth_error; + + // IR-LSQ-mode fields + long ir_total_us; + long ir_setup_us = 0; // warm-start x0 build time, its OWN slot (NOT inside + // ir_total_us; 0 = cold start). Assigned for the + // Blendenpik rows. + int ir_outer_iters; + int ir_inner_iters_total; + int lsqr_iters = 0; // LSQR iterations, where an LSQR phase ran (published + // Blendenpik rows only) + int engine_status = -1; // restarted_pcg_ne exit status (-1 = engine not used) + std::string stop_reason = "n/a"; // named exit condition (see the reason helpers) + T x0_relres = (T)-1; // true relres of the handed-off warm x0 (refine warm rows) + T ls_residual_norm; + T ls_solution_error; // -1 sentinel when undefined (FEM irlsq) + + // Inner-CG diagnosis. ir_inner_iters_total alone cannot say whether a + // solve converged or merely exhausted its budget: both used to report success. + // ir_inner_capped : 1 if ANY outer step hit max_inner (0 = all converged, + // 2 = a CG breakdown occurred). Published + // Blendenpik rows carry 0/1 from LSQR's own convergence, + // refine rows carry real kernel diagnoses; -1 only on + // failed builds. + // ir_inner_relres : worst per-step achieved ||Mz-c||/||c|| (NE space). + // CAVEAT: for published Blendenpik rows this column holds + // LSQR's final LS-space relres instead; do not compare the + // two families through this column without saying so. + // ir_inner_best_relres / ir_inner_best_iter : the smallest residual seen in the + // worst step and the iteration it happened at. best_iter far below the + // iteration count means the solve STAGNATED (tolerance below the attainable + // floor); best_iter near the count means it was still converging when capped. + int ir_inner_capped = -1; + T ir_inner_relres = (T)-1; + T ir_inner_best_relres = (T)-1; + int ir_inner_best_iter = -1; + + // cond(J R^-1): the number that actually says whether the preconditioner works. + // -1 sentinel when not computed: the compute_cond CLI flag now gates this + // computation directly (both irlsq and irlsq_reg), so -1 means + // either compute_cond was off or n exceeded the internal eig cap (16384). + T cond_precond = (T)-1; + + // irlsq_reg only: kappa(A) estimate from the regularized R diagonal + // (max|R_ii| / min|R_ii|; floored near sigma_max/mu when sigma_min < mu). + // -1 sentinel for the plain irlsq path. + T kappa_measured = (T)-1; + + // QR timing breakdown (from algo.times[]) + std::vector qr_breakdown; + std::vector ir_breakdown; + + // Per-round engine records: filled for + // every row that ran the shared engine (IR methods and refine rows); empty for + // published Blendenpik rows and failed builds. Written to the *_rounds.csv sidecar. + std::vector round_iters, round_status, round_best_iter; + std::vector round_relres, round_best_relres, round_ls_relres; + + // RSS WINDOW SEMANTICS, per path: + // irlsq / rspec: Q-less rows stop the tracker right after the QR build + // (build-only peak, matching the build-phase analytical models); + // Blendenpik-family rows stop after run_blendenpik_family returns, so + // their window also spans the LSQR/engine solve. analytical_kb models + // each window kind correspondingly (build-only vs. build-plus-LSQR). + // irlsq_reg: windows are UNIFIED across every algorithm, build + solve, + // diagnostics excluded (the orth-loss materialization is moved outside + // the window for exactly this reason; see the ordering note in + // run_irlsq_reg). NOTE: analytical_kb for the Q-less rows there is still + // assigned from the build-only QR model at dispatch time, before the + // solve runs, so peak_rss_kb (build+solve) and analytical_kb (build-only) + // are no longer measuring the same window for those rows (a known + // mismatch this comment records but does not correct). + long peak_rss_kb; + long analytical_kb; +}; + +// Fold a driver's per-pass shift record into the result row (shared fold in +// cqrrt_bench_common.hh; see record note there). +template +static void record_chol_shift(bench_result& res, const T (&shifts)[N], const T (&traces)[N]) { + fold_chol_shift(res.chol_shift_abs, res.chol_shift_rel, shifts, traces); +} + +// ---- CLI-configurable inner-CG controls ------------------------------------- +// File-scope rather than threaded through the runners, whose signatures already take +// 14 parameters. Both are set once in main() from argv before any runner is called and +// are read-only thereafter. +// +// Why they are configurable at all: the inner-CG budget was hard-coded at +// 200 per outer step, which with g_ir_n_steps outer steps produced a fixed iteration +// ceiling in the CSVs, and the tolerance was fixed at eps^0.85 (~4.9e-14 in double), a +// relative-residual target close enough to the floating-point stagnation floor that CG +// can be unable to reach it regardless of preconditioner quality. Exposing both lets a +// diagnostic sweep separate those two effects without a rebuild. +static int g_ir_max_inner = 200; // <= 0 => keep the IterRefineLSQ default +static double g_ir_inner_tol = -1.0; // < 0 => eps^0.85 in the working precision +static double g_ir_round_drop = 1e-4; // per-round CG drop; 0 = legacy fixed-tol rounds +static int g_ir_n_steps = 50; // outer-round cap (campaign-canonical 50; the + // cap must not bind before tol + maxit do: + // native_ill CholQR2 genuinely uses 50 rounds.) +static double g_ir_outer_tol = -1.0; // <0 => 10*eps(solve precision); 0 disables early exit +// (g_ir_warm_start / g_bp_warm_start are gone: IR methods are always +// cold; Blendenpik runs as two mask-32 variants, warm and cold.) + +// Outer-stagnation window override: RANDLAPACK_IR_OUTER_STAG, read +// once. Default 2 (the engine default); 0 disables the LS-floor exit. The +// knob exists for the CholQR full-precision diagnostic cell, which must show +// the flat trajectory rather than exit at it. Env rather than CLI so the +// campaign arg layout stays frozen. +static int ir_outer_stag_window() { + static const int w = []() { + const char* s = std::getenv("RANDLAPACK_IR_OUTER_STAG"); + if (s == nullptr || *s == '\0') return 2; + char* end = nullptr; + long v = std::strtol(s, &end, 10); + // atoi silently returned 0 on garbage, which silently DISABLES the + // LS-floor exit rather than erroring (0 is also a legal value chosen + // for that purpose, so a typo cannot be told apart from intent). + if (end == s || *end != '\0') { + std::cerr << "FATAL: RANDLAPACK_IR_OUTER_STAG='" << s + << "' is not a valid integer; refusing to silently run " + "with a possibly-disabled LS-floor exit. Fix or unset it.\n"; + std::exit(1); + } + return (int)v; + }(); + return w; +} + +// Full argv, space-joined and double-quoted, set once in run_benchmark() from +// argc/argv; echoed in every results CSV header so a CSV can be +// traced back to the exact invocation that produced it without a side log. +// quote_join_argv itself is shared (cqrrt_bench_common.hh). +static std::string g_argv_line; + +// Environment provenance for every results CSV: the shared env line +// (cqrrt_bench_common.hh) plus this file's own IR-knob echo. +static void write_env_provenance(std::ofstream& out) { + write_host_line(out); + write_env_line(out); + out << "# ir knobs: max_inner=" << g_ir_max_inner << " inner_tol=" << g_ir_inner_tol + << " round_drop=" << g_ir_round_drop << " n_steps=" << g_ir_n_steps + << " outer_tol=" << g_ir_outer_tol + << " outer_stag_window=" << ir_outer_stag_window() << "\n"; +} + +// Summarize an IterRefineLSQ run's inner-CG behavior into the CSV fields. +// +// Reports the WORST outer step, since one capped step is enough to make the reported +// iteration count meaningless as a convergence measure. `capped` is 0 if every step +// converged, 1 if any step exhausted max_inner, 2 if any step broke down. +template +static void record_inner_cg_diagnosis(const RandLAPACK::IterRefineLSQ& ir, + bench_result& res) { + if (ir.inner_status_per_step.empty()) return; + // Rank by SEVERITY, not by the enum's numeric value. The codes are not ordered by + // severity: Stagnated = 3 was added after Breakdown = 2, so a plain `>` comparison would + // let a clean stagnation (which exits early WITH the best iterate) mask a genuine CG + // breakdown in another step. Severity order, worst first: + // Breakdown (2): solver failed outright + // HitCap (1): ran out of budget while still descending + // Stagnated (3): reached its floor and stopped; benign, best iterate returned + // Converged (0): met the tolerance + auto severity = [](int status) -> int { + switch (status) { + case 2: return 3; // Breakdown + case 1: return 2; // HitCap + case 3: return 1; // Stagnated + default: return 0; // Converged + } + }; + int worst = ir.inner_status_per_step[0]; + size_t worst_idx = 0; + for (size_t i = 1; i < ir.inner_status_per_step.size(); ++i) { + if (severity(ir.inner_status_per_step[i]) > severity(worst)) { + worst = ir.inner_status_per_step[i]; + worst_idx = i; + } + } + // All steps converged: report the step that got the least far. + if (worst == 0 && !ir.inner_relres_per_step.empty()) { + for (size_t i = 0; i < ir.inner_relres_per_step.size(); ++i) + if (ir.inner_relres_per_step[i] > ir.inner_relres_per_step[worst_idx]) worst_idx = i; + } + res.ir_inner_capped = worst; + if (worst_idx < ir.inner_relres_per_step.size()) + res.ir_inner_relres = ir.inner_relres_per_step[worst_idx]; + if (worst_idx < ir.inner_best_relres_per_step.size()) + res.ir_inner_best_relres = ir.inner_best_relres_per_step[worst_idx]; + if (worst_idx < ir.inner_best_iter_per_step.size()) + res.ir_inner_best_iter = ir.inner_best_iter_per_step[worst_idx]; +} + +// Same diagnosis, from a raw engine history (the refine rows bypass IterRefineLSQ). +// Identical severity ranking. +template +static void record_inner_cg_diagnosis(const RandLAPACK::PCGRoundHistory& h, + bench_result& res) { + if (h.status.empty()) return; + auto severity = [](int status) -> int { + switch (status) { + case 2: return 3; // Breakdown + case 1: return 2; // HitCap + case 3: return 1; // Stagnated + default: return 0; // Converged + } + }; + size_t worst_idx = 0; + for (size_t i = 1; i < h.status.size(); ++i) + if (severity(h.status[i]) > severity(h.status[worst_idx])) worst_idx = i; + if (h.status[worst_idx] == 0) { + for (size_t i = 0; i < h.relres.size(); ++i) + if (h.relres[i] > h.relres[worst_idx]) worst_idx = i; + } + res.ir_inner_capped = h.status[worst_idx]; + res.ir_inner_relres = h.relres[worst_idx]; + res.ir_inner_best_relres = h.best_relres[worst_idx]; + res.ir_inner_best_iter = h.best_iter[worst_idx]; +} + +// Round-record copiers for the *_rounds.csv sidecar. +template +static void copy_round_records(const RandLAPACK::PCGRoundHistory& h, bench_result& res) { + res.round_iters = h.iters; + res.round_status = h.status; + res.round_best_iter = h.best_iter; + res.round_relres = h.relres; + res.round_best_relres = h.best_relres; + res.round_ls_relres = h.ls_relres; +} +template +static void record_ir_outputs(const RandLAPACK::IterRefineLSQ& ir, bench_result& res) { + res.engine_status = ir.engine_status; + res.stop_reason = pcg_stop_reason(ir.engine_status); + res.round_iters = ir.inner_iters_per_step; + res.round_status = ir.inner_status_per_step; + res.round_best_iter = ir.inner_best_iter_per_step; + res.round_relres = ir.inner_relres_per_step; + res.round_best_relres = ir.inner_best_relres_per_step; + res.round_ls_relres = ir.ls_relres_per_step; +} + +// Shared method-mask decode, to avoid per-path copies diverging (the rspec +// copy once silently ignored bits 32/64; the console echo once showed bits +// 0-4 only). with_blendenpik = false (rspec) warns on 32/64 instead. +static std::vector decode_method_mask(int64_t method_mask, bool with_blendenpik) { + std::vector algs; + if (method_mask & 1) algs.push_back("CQRRT_linop"); + if (method_mask & 2) algs.push_back("CholQR"); + if (method_mask & 4) algs.push_back("sCholQR3"); + if (method_mask & 8) algs.push_back("sCholQR3_basic"); + if (method_mask & 16) algs.push_back("CholQR2"); + if (with_blendenpik) { + if (method_mask & 32) { // published: its own sketch-and-solve warm start + cold + algs.push_back("Blendenpik"); + algs.push_back("Blendenpik_cold"); + } + if (method_mask & 64) { // refined by the shared engine (refined_blendenpik.hh) + algs.push_back("Blendenpik_refine"); + algs.push_back("Blendenpik_cold_refine"); + } + } else if (method_mask & (32 | 64)) { + std::cerr << "Warning: method_mask bits 32/64 (Blendenpik families) are not " + "available in this mode and are ignored.\n"; + } + return algs; +} + +// ============================================================================ +// Shared Blendenpik-family dispatch (published + refined rows), used by BOTH the +// irlsq path (run_benchmark_inner) and the irlsq_reg path. Extracted after the +// two per-path copies drifted into different wrong accountings: the warm refine +// row silently ran cold in both paths (the exact-string warm_start test was +// false for "Blendenpik_refine"), and the refinement phase was missing from the +// time columns in both, from the iteration count in one. +// ============================================================================ +template +static void run_blendenpik_family( + GLO& A_op, const T* b, int64_t m, T* x_ls, int64_t n, T* R_T, + T d_factor, int64_t sketch_nnz, RandBLAS::RNGState state, + const std::string& alg_name, T tol, T outer_tol_eff, + RandLAPACK::PeakRSSTracker& mem, bench_result& res) +{ + const bool is_ref = (alg_name.find("_refine") != std::string::npos); + const bool warm = (alg_name == "Blendenpik") || (alg_name == "Blendenpik_refine"); + const int max_inner = (g_ir_max_inner > 0) ? g_ir_max_inner : 200; + const int budget = max_inner * g_ir_n_steps; // same budget the IR methods get + const T inner_tol = (g_ir_inner_tol > 0) ? (T)g_ir_inner_tol + : std::pow(std::numeric_limits::epsilon(), (T)0.85); + + if (is_ref) { + // Refined rows: init_only sketch-and-solve x0, ALL iterative work in the + // shared engine, no internal LSQR (see benchmark/refined_blendenpik.hh). + // Engine knobs mirror what IterRefineLSQ passes for the Q-less rows. + const bool paced = (g_ir_round_drop > 0); + const T drop = paced ? (T)g_ir_round_drop : inner_tol; + const T abs_guard = paced ? inner_tol : (T)0; + std::fill(x_ls, x_ls + n, (T)0); + auto rr = RandLAPACK::bench::run_refined_blendenpik( + A_op, b, m, x_ls, n, d_factor, sketch_nnz, state, warm, + outer_tol_eff, budget, max_inner, drop, g_ir_n_steps - 1, + /*stag_window=*/20, /*stag_rel_improve=*/(T)1e-3, abs_guard, + ir_outer_stag_window()); + res.qr_status = rr.qr_status; + res.peak_rss_kb = mem.stop(); + if (res.qr_status != 0) return; + res.qr_time_us = rr.qr_us; + res.ir_setup_us = rr.setup_us; // x0 build (0 for the cold row) + res.x0_relres = rr.x0_relres; // warm-start quality (-1 cold) + std::copy(rr.R, rr.R + rr.R_sz, R_T); + // QR-breakdown slots for Blendenpik-family rows (see the breakdown + // writer header for the full per-algorithm table): t1=qr, t4=x0 setup, + // everything else 0. Refine rows only expose the COMBINED sketch+QR + // time (rr.qr_us is not split further), so it lands whole in t1 and t0 + // (sketch) stays 0, unlike the published branch below, which does + // have the split. + res.qr_breakdown.assign(5, 0L); + res.qr_breakdown[1] = rr.qr_us; + res.qr_breakdown[4] = rr.setup_us; + // Refined rows use init_only, so LSQR never runs; the engine workspace + // that follows (9n + m) is smaller than the sketch term, so the + // Blendenpik moment is still the peak. + res.analytical_kb = RandLAPACK::blendenpik_linops_analytical_kb( + m, n, (double)d_factor, /*warm_start=*/warm, /*with_lsqr=*/false); + res.ir_total_us = rr.solve_us; // the WHOLE refinement solve (was missing) + res.ir_outer_iters = rr.rounds; // real round count (was clobbered to 1) + res.ir_inner_iters_total = rr.iters; // engine inner CG only: single unit + res.lsqr_iters = 0; // no LSQR phase in the redesigned rows + res.engine_status = rr.status; + res.stop_reason = pcg_stop_reason(rr.status); + record_inner_cg_diagnosis(rr.history, res); + copy_round_records(rr.history, res); + // ir_breakdown in the IterRefineLSQ layout [total, inner_cg, trsm, fwd, adj, other]. + long op_outer = (rr.t_fwd_us - rr.history.t_fwd_inner_us) + + (rr.t_adj_us - rr.history.t_adj_inner_us) + + (rr.t_trsm_us - rr.history.t_trsm_inner_us); + long other = rr.solve_us - rr.history.t_inner_us - op_outer; + if (other < 0) other = 0; + res.ir_breakdown = {rr.solve_us, rr.history.t_inner_us, rr.t_trsm_us, + rr.t_fwd_us, rr.t_adj_us, other}; + return; + } + + // Published rows: sketch + QR + LSQR (warm = its own sketch-and-solve x0). + RandLAPACK::Blendenpik_linops bp(/*time_subroutines=*/true, tol); + bp.nnz = sketch_nnz; + bp.warm_start = warm; + bp.max_iters = budget; + std::fill(x_ls, x_ls + n, (T)0); + res.qr_status = bp.call(A_op, b, m, x_ls, n, d_factor, state); + res.peak_rss_kb = mem.stop(); + if (res.qr_status != 0) return; + res.qr_time_us = bp.times[0] + bp.times[1]; // sketch + QR + res.ir_setup_us = warm ? bp.times[4] : 0; // warm x0 build, its own column + std::copy(bp.R_out, bp.R_out + bp.R_out_sz, R_T); + // QR-breakdown slots for Blendenpik-family rows (see the breakdown writer + // header): t0=sketch, t1=qr, t4=x0 setup (0 for the cold row, since x0/Sb/r0 + // are only allocated when warm_start||init_only); everything else 0. + res.qr_breakdown.assign(5, 0L); + res.qr_breakdown[0] = bp.times[0]; + res.qr_breakdown[1] = bp.times[1]; + res.qr_breakdown[4] = bp.times[4]; + // Published rows run LSQR, so the LSQR workspace is live at the peak. + res.analytical_kb = RandLAPACK::blendenpik_linops_analytical_kb( + m, n, (double)d_factor, /*warm_start=*/warm, /*with_lsqr=*/true); + // No IR loop: reuse the diagnosis columns from LSQR's own convergence signals. + res.ir_inner_capped = bp.converged ? 0 : 1; + res.ir_inner_relres = bp.final_relres; + res.ir_total_us = bp.times[2]; + res.ir_outer_iters = 1; + res.ir_inner_iters_total = bp.lsqr_iters; // LSQR iters in the CG-iters slot (published rows only) + res.lsqr_iters = bp.lsqr_iters; + res.engine_status = -1; // the pcg engine did not run + res.stop_reason = lsqr_stop_reason(bp.converged, bp.lsqr_stop_test); + // ir_breakdown from LSQR's op split. + if (bp.lsqr_op_times.size() >= 3) { + long fwd = bp.lsqr_op_times[0], adj = bp.lsqr_op_times[1], trsm = bp.lsqr_op_times[2]; + long other = bp.times[2] - (fwd + adj + trsm); + if (other < 0) other = 0; + res.ir_breakdown = {bp.times[2], 0, trsm, fwd, adj, other}; // no kernel split in LSQR + } +} + +// Per-round sidecar writer: one row per (algorithm, run, round) for +// every row that ran the shared engine. +template +static void write_rounds_csv(const std::string& filename, + const std::vector>& results) { + std::ofstream out(filename); + out << "# Per-round engine records (restarted_pcg_ne / IterRefineLSQ).\n" + << "# inner_status: 0 Converged, 1 HitCap, 2 Breakdown, 3 Stagnated.\n" + << kRoundsCsvHeader; + for (const auto& r : results) { + for (size_t k = 0; k < r.round_iters.size(); ++k) { + write_round_row(out, r.alg_name, r.run_idx, k + 1, + r.round_iters[k], r.round_status[k], r.round_relres[k], + r.round_best_relres[k], r.round_best_iter[k], r.round_ls_relres[k]); + } + } +} + +// estimate_op_2norm and compute_orth_error_explicit are shared with +// bench_toeplitz_ls; see cqrrt_bench_common.hh (using-declared above). + +// ============================================================================ +// CSV writers: IR-LSQ (preserves the column order plot_irlsq_results.m expects) +// ============================================================================ + +template +static void write_irlsq_results( + const std::string& filename, + const std::vector>& results, + int64_t m, int64_t n, int64_t nnz_or_zero, const std::string& input_label, + T noise_level, T d_factor, int64_t sketch_nnz, int64_t block_size, + int64_t method_mask, int64_t num_runs, long chol_time_us, + const std::string& precision_str) +{ + std::ofstream out(filename); + out << "# IR-LSQ Benchmark results\n" + << "# Date: " << make_run_timestamp() << "\n" + << "# argv=" << g_argv_line << "\n" + << "# input=" << input_label << "\n" + << "# precision=" << precision_str << "\n" + << "# M=" << m << " N=" << n << " nnz=" << nnz_or_zero << "\n" + << "# noise_level=" << noise_level << "\n" + << "# chol_time_us=" << chol_time_us << "\n" + << "# d_factor=" << d_factor << " sketch_nnz=" << sketch_nnz + << " block_size=" << block_size << "\n" + << "# method_mask=" << method_mask << "\n" + << "# num_runs=" << num_runs << "\n" +#ifdef _OPENMP + << "# OpenMP threads: " << omp_get_max_threads() << "\n" +#else + << "# OpenMP threads: 1\n" +#endif + ; + write_env_provenance(out); + out << "algorithm,run,m,n,qr_status,qr_time_us,peak_rss_kb,analytical_kb," + "orth_error," + "ir_total_us,ir_outer_iters,ir_inner_iters_total," + "ls_residual_norm,ls_solution_error," + "ir_inner_capped,ir_inner_relres,ir_inner_best_relres,ir_inner_best_iter,cond_precond," + "ir_setup_us,lsqr_iters,engine_status,stop_reason,x0_relres,chol_shift_abs,chol_shift_rel\n"; + // Sentinel note: chol_shift_abs/chol_shift_rel use -1 for "no Cholesky in + // this row" (Blendenpik family, unpreconditioned) or "QR failed before a + // shift record existed"; 0 still means "Cholesky ran unshifted". + for (const auto& r : results) { + out << r.alg_name << "," << r.run_idx << "," << r.m << "," << r.n << "," + << r.qr_status << "," << r.qr_time_us << "," << r.peak_rss_kb << "," << r.analytical_kb << "," + << std::scientific << std::setprecision(6) << r.orth_error << "," + << r.ir_total_us << "," << r.ir_outer_iters << "," << r.ir_inner_iters_total << "," + << std::scientific << std::setprecision(6) << r.ls_residual_norm << "," + << std::scientific << std::setprecision(6) << r.ls_solution_error << "," + << r.ir_inner_capped << "," + << std::scientific << std::setprecision(6) << r.ir_inner_relres << "," + << std::scientific << std::setprecision(6) << r.ir_inner_best_relres << "," + << r.ir_inner_best_iter << "," + << std::scientific << std::setprecision(6) << r.cond_precond << "," + << r.ir_setup_us << "," << r.lsqr_iters << "," << r.engine_status << "," + << r.stop_reason << "," + << std::scientific << std::setprecision(6) << r.x0_relres << "," + << std::scientific << std::setprecision(6) << r.chol_shift_abs << "," + << std::scientific << std::setprecision(6) << r.chol_shift_rel + << "\n"; + } +} + +// Write one breakdown row: pads/truncates the phase vector to exactly 18 +// columns (t0..t17: sCholQR3's 18 slots and sCholQR3_basic's 15 are not cut to +// 11) and appends `total_val` as a dedicated final +// t_total column so phase bars can be validated against the row's own +// authoritative total (qr_time_us / ir_total_us) without inferring it from +// vector position, which used to differ silently per algorithm. +static void write_breakdown_row(std::ofstream& out, const std::string& alg, + int64_t run_idx, const char* phase, + const std::vector& v, long total_val) { + out << alg << "," << run_idx << "," << phase; + for (int i = 0; i < 18; ++i) + out << "," << (i < (int)v.size() ? v[i] : 0L); + out << "," << total_val << "\n"; +} + +template +static void write_irlsq_breakdown( + const std::string& filename, + const std::vector>& results, + const std::string& mode_label) +{ + std::ofstream out(filename); + out << "# " << mode_label << " Benchmark runtime breakdown (microseconds)\n" + << "# QR breakdown layout depends on algorithm:\n" + << "# CQRRT_linop (t0-t10): alloc,saso,qr,precond_inv,fwd,adj,gemm,chol,finalize,rest,total\n" + << "# CholQR (t0-t5): alloc,fwd,adj,chol,rest,total (t6-t17 = 0)\n" + << "# CholQR2 (t0-t10): alloc,fwd1,adj1,chol1,upd1,fwd2,adj2,gemm2,chol2,upd2,total (t11-t17 = 0)\n" + << "# sCholQR3_basic(t0-t14): alloc,fwd1,adj1,chol1,trsm1=0,fwd_q=0,syrk2,chol2,upd2,\n" + << "# syrk3,chol3,upd3,q_mat,rest,total (t15-t17 = 0)\n" + << "# sCholQR3 (t0-t17): alloc,fwd1,adj1,chol1,upd1,fwd2,adj2,gemm2,chol2,upd2,\n" + << "# fwd3,adj3,gemm3,chol3,upd3,q_mat,rest,total\n" + << "# Blendenpik-family (t0-t4 only, t5-t17 = 0): t0=sketch, t1=qr, t4=x0 setup (0 for\n" + << "# cold rows). Refine rows only expose the combined sketch+QR time, so it lands\n" + << "# whole in t1 with t0 left 0.\n" + << "# t_total is the row's own authoritative total (qr_time_us for the QR phase row,\n" + << "# ir_total_us for the IR phase row), independent of how far t0..t17 are populated.\n" + << "# IR-LSQ breakdown (6, in t0-t5; t6-t17 = 0): outer_total, inner_cg_total, trsm_total,\n" + << "# fwd_total, adj_total, other\n" + << "# (t_total on the IR row equals ir_total_us; Blendenpik rows carry\n" + << "# real IR entries too: refine rows the engine split, published rows the LSQR op\n" + << "# split with a 0 inner_cg slot)\n" + << "algorithm,run,phase,t0,t1,t2,t3,t4,t5,t6,t7,t8,t9,t10,t11,t12,t13,t14,t15,t16,t17,t_total\n"; + for (const auto& r : results) { + write_breakdown_row(out, r.alg_name, r.run_idx, "QR", r.qr_breakdown, r.qr_time_us); + write_breakdown_row(out, r.alg_name, r.run_idx, "IR", r.ir_breakdown, r.ir_total_us); + } +} + +// ============================================================================ +// CSV writer: RSPEC (reduced spectral approximation) +// ============================================================================ + +template +struct rspec_result { + int64_t m; + int64_t n; + int64_t run_idx; + std::string alg_name; + int qr_status; + long qr_time_us; + long peak_rss_kb; + long analytical_kb; + long factor_time_us; + long rspec_total_us; + T orth_error; // ||Q^T Q - I||_F / sqrt(n), Q = V_app R^{-1} + std::vector qr_breakdown; // Q-less QR breakdown (driver times[], same layout as irlsq) + std::vector rr_breakdown; // Rayleigh-Ritz post-processing: [orth, rr_build, syevd, resid] (us) + std::vector top_eigvals; + std::vector top_residuals; +}; + +template +static void write_rspec_breakdown( + const std::string& filename, + const std::vector>& results) +{ + std::ofstream out(filename); + out << "# RSPEC runtime breakdown (microseconds)\n" + << "# QR breakdown layout depends on algorithm; see write_irlsq_breakdown's header\n" + << "# in CQRRT_linop_applications.cc for the full per-algorithm slot table.\n" + << "# RR breakdown (4, in t0-t3; t4-t17 = 0): orth_error, rayleigh_ritz_build, syevd, ritz_residuals\n" + << "# t_total on the QR row is qr_time_us; on the RR row it is the sum of t0..t3 (rspec has\n" + << "# no separately tracked RR total field; rspec_total_us excludes the orth diagnostic,\n" + << "# see the file header comment on why).\n" + << "algorithm,run,phase,t0,t1,t2,t3,t4,t5,t6,t7,t8,t9,t10,t11,t12,t13,t14,t15,t16,t17,t_total\n"; + for (const auto& r : results) { + write_breakdown_row(out, r.alg_name, r.run_idx, "QR", r.qr_breakdown, r.qr_time_us); + // -1 sentinel on a failed QR row (rr_breakdown is cleared, and a sum + // over an empty vector is 0, which would misreport as "measured, zero + // cost" rather than "not run"). + long rr_total = -1; + if (r.qr_status == 0) { + rr_total = 0; + for (long v : r.rr_breakdown) rr_total += v; + } + write_breakdown_row(out, r.alg_name, r.run_idx, "RR", r.rr_breakdown, rr_total); + } +} + +template +static void write_rspec_csv( + const std::string& filename, + const std::vector>& results, + int64_t m, int64_t n, int num_runs, + const std::string& K_file, const std::string& M_file, const std::string& V_file, + double omega, int64_t power_j, + int64_t sketch_nnz, int64_t block_size, + int64_t method_mask, int top_k) +{ + std::ofstream out(filename); + out << "# RSPEC (reduced spectral approximation) results\n" + << "# Date: " << make_run_timestamp() << "\n"; + write_host_line(out); + out << "# argv=" << g_argv_line << "\n" + << "# Matrix dimensions: m=" << m << " n=" << n << "\n" + << "# Runs per algorithm: " << num_runs << "\n" +#ifdef _OPENMP + << "# OpenMP threads: " << omp_get_max_threads() << "\n" +#else + << "# OpenMP threads: 1\n" +#endif + << "# K_file: " << K_file << "\n" + << "# M_file: " << M_file << "\n" + << "# V_file: " << V_file << "\n" + << "# omega: " << omega << "\n" + << "# power_j: " << power_j << "\n" + << "# sketch_nnz: " << sketch_nnz << "\n" + << "# block_size: " << block_size << "\n" + << "# method_mask: " << method_mask << "\n" + << "# top_k: " << top_k << "\n"; + + out << "algorithm,run,m,n,omega,power_j,qr_status,qr_time_us,peak_rss_kb,analytical_kb," + "factor_time_us,rspec_total_us,orth_error"; + for (int i = 0; i < top_k; ++i) out << ",eig_" << i; + for (int i = 0; i < top_k; ++i) out << ",resid_" << i; + out << "\n"; + + for (const auto& r : results) { + out << r.alg_name << "," << r.run_idx << "," << r.m << "," << r.n << "," + << omega << "," << power_j << "," + << r.qr_status << "," << r.qr_time_us << "," + << r.peak_rss_kb << "," << r.analytical_kb << "," + << r.factor_time_us << "," << r.rspec_total_us << "," + << std::scientific << std::setprecision(6) << r.orth_error; + for (int i = 0; i < top_k; ++i) { + T v = (i < (int)r.top_eigvals.size()) ? r.top_eigvals[i] + : std::numeric_limits::quiet_NaN(); + out << "," << std::scientific << std::setprecision(8) << v; + } + for (int i = 0; i < top_k; ++i) { + T v = (i < (int)r.top_residuals.size()) ? r.top_residuals[i] + : std::numeric_limits::quiet_NaN(); + out << "," << std::scientific << std::setprecision(6) << v; + } + out << "\n"; + } +} + +// ============================================================================ +// Console summary +// ============================================================================ + +template +static void print_irlsq_summary(const bench_result& r) { + std::printf("\n [%s] Run %lld (noise=%.3f):\n", + r.alg_name.c_str(), (long long)r.run_idx, (double)r.noise_level); + if (r.qr_status != 0) { + std::printf(" QR returned status %d, IR-LSQ skipped.\n", r.qr_status); + return; + } + std::printf(" QR: %lld us, peak_RSS=%lld KB, predicted=%lld KB\n", + (long long)r.qr_time_us, (long long)r.peak_rss_kb, (long long)r.analytical_kb); + if (r.orth_error >= 0) std::printf(" orth_err = %.3e\n", (double)r.orth_error); + std::printf(" IR-LSQ (x_0=0): total=%lld us, outer=%d, inner_total=%d\n", + (long long)r.ir_total_us, r.ir_outer_iters, r.ir_inner_iters_total); + std::printf(" ||Ax-b||/(||A||*||x||+||b||) = %.3e\n", (double)r.ls_residual_norm); + if (r.ls_solution_error >= 0) + std::printf(" ||x-x_true||/||x_true|| = %.3e\n", (double)r.ls_solution_error); + else + std::printf(" ||x-x_true||/||x_true|| = N/A (no ground-truth x_true)\n"); +} + +// ============================================================================ +// Core templated runner +// ============================================================================ + +template +static int run_benchmark_inner( + OpType& A_op, + int64_t m, int64_t n, int64_t input_nnz, + const std::string& output_dir, int64_t num_runs, + T d_factor, int64_t sketch_nnz, int64_t block_size, + bool compute_cond, + int64_t method_mask, T noise_level, + long chol_time_us, + const std::string& op_label, + const std::string& input_label, + const std::vector* b_ptr, // M-vector RHS + const std::vector* x_true_ptr) // N-vector ground truth (sparse only); nullptr otherwise +{ + // b_ptr used to be treated as optional (see the now-removed `if (b_ptr)` + // guards below), but every call site always passes a real RHS and the + // post-processing block dereferences it unconditionally regardless; make + // that assumption explicit instead of leaving a misleading nullable API. + randlapack_require(b_ptr != nullptr) + << "run_benchmark_inner: b_ptr (the RHS vector) must be non-null."; + + // Ordered list of selected algorithm names from the bitmask (shared decode; + // see decode_method_mask and the mask documentation in the file header). + std::vector selected_algs = decode_method_mask(method_mask, /*with_blendenpik=*/true); + + if (selected_algs.empty()) { + std::cerr << "Error: method_mask selects no algorithms (got " << method_mask << ").\n"; + return 1; + } + + if (compute_cond) { + RandLAPACK::testing::print_condition_diagnostics(A_op, op_label); + } + + // Per-run RNG states + RandBLAS::RNGState main_state(123); + std::vector> run_states(num_runs); + for (int64_t r = 0; r < num_runs; ++r) { + run_states[r] = main_state; + if (r > 0) run_states[r].key.incr(r); + } + + T tol = std::pow(std::numeric_limits::epsilon(), (T)0.85); + + // Warmup (CQRRT_linop), plus the solve path: the build warmup already + // applies A_op repeatedly, but the timed IterRefineLSQ also exercises the + // TRSM preconditioner path and LSQR's vector work, whose one-time costs + // (thread pools, first-touch pages) otherwise land inside the FIRST + // method's timed solve. CPU warmup only, + // distinct from the x0 warm-start ablation. + std::cout << "Running warmup... " << std::flush; + { + auto warm_state = run_states[0]; + T* R_warm = new T[n * n](); + RandLAPACK::CQRRT_linops warm_algo(false, tol, false); + warm_algo.nnz = sketch_nnz; + warm_algo.block_size = block_size; + int warm_status = warm_algo.call(A_op, R_warm, n, d_factor, warm_state); + T* x_wu = new T[n](); + int it_wu = 0; long lt_wu[4] = {0}; + RandLAPACK::lsqr(A_op, m, n, + (warm_status == 0) ? R_warm : nullptr, + (warm_status == 0) ? n : (int64_t)0, + b_ptr->data(), x_wu, tol, tol, 5, it_wu, lt_wu); + delete[] x_wu; + delete[] R_warm; + } + std::cout << "done\n\n"; + + // Precompute ||A||_2 and ||b|| for the Higham backward-error metric: + // ls_residual_norm = ||A x - b|| / (||A||_2 * ||x|| + ||b||) + std::cout << "Estimating ||A||_2 via power iteration (10 iters)... " << std::flush; + T A_2norm = estimate_op_2norm(A_op, m, n, 10); + T b_norm = blas::nrm2(m, b_ptr->data(), 1); + std::cout << "||A||_2 ~ " << A_2norm << ", ||b|| = " << b_norm << "\n\n"; + + T x_true_norm = (T)0; + if (x_true_ptr) x_true_norm = blas::nrm2(n, x_true_ptr->data(), 1); + + std::vector> all_results; + + // Per-iteration workspaces, hoisted once: invariant sizes across all (alg, run) iters. + T* R = new T[n * n](); // QR output; zero-filled per iter to match prior behavior + T* x_ls = new T[n]; // initial guess (x_0 = 0) + refined solution + T* Ax = new T[m]; // A * x_ls for residual; overwritten beta=0 + + // ================================================================ + // Per-(method, run) loop + // ================================================================ + for (const auto& alg_name : selected_algs) { + std::cout << "\n=== Algorithm: " << alg_name << " ===\n"; + + for (int64_t run_idx = 0; run_idx < num_runs; ++run_idx) { + bench_result res{}; + res.m = m; res.n = n; + res.run_idx = run_idx; + res.alg_name = alg_name; + res.noise_level = noise_level; + res.chol_time_us = chol_time_us; + res.qr_status = 0; + res.qr_time_us = 0; + res.orth_error = (T)-1.0; + res.ir_total_us = 0; + res.ir_outer_iters = 0; + res.ir_inner_iters_total = 0; + res.ls_residual_norm = (T)-1.0; + res.ls_solution_error = (T)-1.0; + res.peak_rss_kb = 0; + res.analytical_kb = -1; // -1 = no value; 0 would read as a real 0 MB bar + + std::fill(R, R + n * n, (T)0); + auto state = run_states[run_idx]; + + // ---- QR dispatch (lifted verbatim from CQRRT_linop_irlsq.cc; +Blendenpik) ---- + std::cout << "[Run " << run_idx << ", " << alg_name << "] QR ... " << std::flush; + RandLAPACK::PeakRSSTracker mem; mem.start(); + if (alg_name.rfind("Blendenpik", 0) == 0) { + // Shared Blendenpik-family dispatch (published + refined rows); fills + // every accounting field itself, see run_blendenpik_family. + T outer_tol_eff = (g_ir_outer_tol >= 0) ? (T)g_ir_outer_tol + : (T)10 * std::numeric_limits::epsilon(); + run_blendenpik_family(A_op, b_ptr->data(), m, x_ls, n, R, + d_factor, sketch_nnz, state, alg_name, tol, outer_tol_eff, mem, res); + } else if (alg_name == "sCholQR3") { + RandLAPACK::sCholQR3_linops qr_algo(/*time_subroutines=*/true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + qr_algo.block_size = block_size; + res.qr_status = qr_algo.call(A_op, R, n); + record_chol_shift(res, qr_algo.chol_applied_shifts, qr_algo.chol_gram_traces); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::scholqr3_linops_analytical_kb(m, n, block_size); + } + } else if (alg_name == "sCholQR3_basic") { + RandLAPACK::sCholQR3_linops_basic qr_algo(/*time_subroutines=*/true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + res.qr_status = qr_algo.call(A_op, R, n); + record_chol_shift(res, qr_algo.chol_applied_shifts, qr_algo.chol_gram_traces); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::scholqr3_linops_basic_analytical_kb(m, n); + } + } else if (alg_name == "CholQR") { + RandLAPACK::CholQR_linops qr_algo(/*time_subroutines=*/true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + qr_algo.block_size = block_size; + res.qr_status = qr_algo.call(A_op, R, n); + record_chol_shift(res, qr_algo.chol_applied_shifts, qr_algo.chol_gram_traces); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector (6 entries; writer pads) + res.analytical_kb = RandLAPACK::cholqr_linops_analytical_kb(m, n, block_size); + } + } else if (alg_name == "CholQR2") { + RandLAPACK::CholQR2_linops qr_algo(/*time_subroutines=*/true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + qr_algo.block_size = block_size; + res.qr_status = qr_algo.call(A_op, R, n); + record_chol_shift(res, qr_algo.chol_applied_shifts, qr_algo.chol_gram_traces); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::cholqr2_linops_analytical_kb(m, n, block_size); + } + } else { + // CQRRT_linop (TRSM_IDENTITY precond). CQRRT_linop_bqrrp is not + // part of the benchmark dispatch. + RandLAPACK::CQRRT_linops qr_algo(/*time_subroutines=*/true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + qr_algo.nnz = sketch_nnz; + qr_algo.block_size = block_size; + qr_algo.precond_method = RandLAPACK::CQRRTLinopPrecond::TRSM_IDENTITY; + res.qr_status = qr_algo.call(A_op, R, n, d_factor, state); + record_chol_shift(res, qr_algo.chol_applied_shifts, qr_algo.chol_gram_traces); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::cqrrt_linops_analytical_kb(m, n, d_factor, block_size); + } + } + + if (res.qr_status != 0) { + std::cerr << "\n [" << alg_name << "] Run " << run_idx + << ": QR returned status " << res.qr_status + << " (likely Cholesky breakdown). Skipping post-processing.\n"; + res.qr_time_us = -1; + res.ir_total_us = -1; // -1 sentinel: breakdown's t_total must not read as 0 (real cost) + res.qr_breakdown.clear(); // writer pads to 18 zero columns regardless of size + res.analytical_kb = -1; // -1 = no value; 0 would read as a real 0 MB bar + all_results.push_back(res); + print_irlsq_summary(res); + continue; + } + std::cout << "done (" << res.qr_time_us << " us)"; + + // ---- Orth_error: ||Q^T Q - I||_F / sqrt(n), blocked compute. Runs for every method. ---- + // compute_cond means exactly one thing everywhere: whether cond_precond + // gets computed at all (subject to the n<=16384 eig cap inside + // compute_orth_error_explicit). + res.orth_error = compute_orth_error_explicit(A_op, R, m, n, block_size, + compute_cond ? &res.cond_precond : nullptr); + + // ---- IR-LSQ post-processing ---- + { + const std::vector& b = *b_ptr; + if (alg_name.rfind("Blendenpik", 0) == 0) { + // x_ls and every solve-accounting field were already produced by + // run_blendenpik_family; nothing to overwrite here (overwriting + // ir_total_us / ir_outer_iters / ir_inner_iters_total with + // LSQR-only values here would erase the refinement). + std::cout << ". solve recorded ... " << std::flush; + } else { + std::cout << ". IR-LSQ ... " << std::flush; + auto ls_t0 = steady_clock::now(); + + // Initial guess x_0 = 0 (per collaborator: no sketching in the LS + // solve itself). The only randomness is S_1 inside Q-less QR, which + // yields the preconditioner R; IterRefineLSQ starts from zero and the + // preconditioned inner CG converges from there. + std::fill(x_ls, x_ls + n, (T)0.0); + + RandLAPACK::IterRefineLSQ ir( + /*tol=*/ (g_ir_inner_tol > 0) ? (T)g_ir_inner_tol : tol, + /*max_inner=*/(g_ir_max_inner > 0) ? g_ir_max_inner : 200, + /*n_steps=*/g_ir_n_steps, + /*timing=*/true, + /*verbose=*/false); + ir.round_drop = (T)g_ir_round_drop; + ir.outer_tol = (g_ir_outer_tol >= 0) ? (T)g_ir_outer_tol + : (T)10 * std::numeric_limits::epsilon(); + ir.outer_stag_window = ir_outer_stag_window(); + int ir_status = ir.call(A_op, R, n, b.data(), m, x_ls, n); + auto ls_t1 = steady_clock::now(); + if (ir_status != 0) { + std::cerr << "Warning: IterRefineLSQ status " << ir_status << " (CG breakdown)\n"; + } + + res.ir_total_us = duration_cast(ls_t1 - ls_t0).count(); + res.ir_outer_iters = ir.outer_iters_done; + res.ir_inner_iters_total = 0; + for (int v : ir.inner_iters_per_step) res.ir_inner_iters_total += v; + record_inner_cg_diagnosis(ir, res); + record_ir_outputs(ir, res); + if (!ir.times.empty()) res.ir_breakdown = ir.times; + } + + // Higham normwise backward-error metric: + // ls_residual_norm = ||A x - b|| / (||A||_2 * ||x|| + ||b||) + // Drivable to machine epsilon for a backward-stable LS solver. + A_op(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, x_ls, n, (T)0.0, Ax, m); + T resid_sq = 0; + #pragma omp parallel for reduction(+:resid_sq) schedule(static) + for (int64_t i = 0; i < m; ++i) { T d = Ax[i] - b[i]; resid_sq += d * d; } + T resid_norm = std::sqrt(resid_sq); + T x_norm = blas::nrm2(n, x_ls, 1); + T denom = A_2norm * x_norm + b_norm; + res.ls_residual_norm = (denom > 0) ? resid_norm / denom : (T)-1.0; + + if (x_true_ptr) { + T err_sq = 0; + for (int64_t i = 0; i < n; ++i) { + T d = x_ls[i] - (*x_true_ptr)[i]; + err_sq += d * d; + } + res.ls_solution_error = (x_true_norm > 0) ? std::sqrt(err_sq) / x_true_norm : (T)-1.0; + } else { + res.ls_solution_error = (T)-1.0; + } + std::cout << "done (" << res.ir_total_us << " us)\n"; + } + + print_irlsq_summary(res); + all_results.push_back(res); + } + } + + // ================================================================ + // CSV output + // ================================================================ + std::string time_buf = make_run_timestamp(); + + std::string results_file = output_dir + "/" + time_buf + "_irlsq_results.csv"; + std::string breakdown_file = output_dir + "/" + time_buf + "_irlsq_breakdown.csv"; + std::string rounds_file = output_dir + "/" + time_buf + "_irlsq_rounds.csv"; + // x_true_ptr is passed only in sparse mode (see the parameter comment above); + // FEM irlsq has no ground truth. Used here purely to label the CSVs: the + // "Sparse IR-LSQ" title used to be hard-coded for the FEM rows too. + const bool is_sparse_input = (x_true_ptr != nullptr); + const std::string mode_label = is_sparse_input ? "Sparse IR-LSQ" : "FEM IR-LSQ"; + const std::string precision_str = (sizeof(T) == 8) ? "double" : "single"; + write_irlsq_results(results_file, all_results, m, n, input_nnz, input_label, + noise_level, d_factor, sketch_nnz, block_size, method_mask, + num_runs, chol_time_us, precision_str); + std::cout << "\nIR-LSQ results written to " << results_file << "\n"; + write_irlsq_breakdown(breakdown_file, all_results, mode_label); + std::cout << "IR-LSQ breakdown written to " << breakdown_file << "\n"; + write_rounds_csv(rounds_file, all_results); + std::cout << "IR-LSQ per-round records written to " << rounds_file << "\n"; + + delete[] R; delete[] x_ls; delete[] Ax; + return 0; +} + +// ============================================================================ +// RSPEC mode runner +// ============================================================================ + +template +static int run_rspec_benchmark( + VAppOpType& V_app_op, // m_K x n_V composite: C^j * V_FEM + CompCOp& C_op, // m_K x m_K composite: L^T X^{-1} L (the operator we Rayleigh-Ritz) + int64_t m_K, int64_t n_V, + const std::string& output_dir, int64_t num_runs, + T d_factor, int64_t sketch_nnz, int64_t block_size, + int64_t method_mask, + long factor_time_us, + const std::string& K_file, const std::string& M_file, const std::string& V_file, + double omega, int64_t power_j) +{ + // Shared decode; rspec has no Blendenpik rows, so bits 32/64 now WARN instead + // of being silently ignored (the old per-path copy's behavior). + std::vector selected_algs = decode_method_mask(method_mask, /*with_blendenpik=*/false); + + if (selected_algs.empty()) { + std::cerr << "Error: method_mask selects no algorithms (got " << method_mask << ").\n"; + return 1; + } + + int64_t m = m_K; + int64_t n = n_V; + int top_k = (int)std::min(10, n); + + // Per-run RNG states (same scheme as run_benchmark_inner). + RandBLAS::RNGState main_state(123); + std::vector> run_states(num_runs); + for (int64_t r = 0; r < num_runs; ++r) { + run_states[r] = main_state; + if (r > 0) run_states[r].key.incr(r); + } + + T tol = std::pow(std::numeric_limits::epsilon(), (T)0.85); + + // Warmup so the Cholesky-factored X^{-1} chain inside V_app_op is warm. + std::cout << "Running rspec warmup... " << std::flush; + { + auto warm_state = run_states[0]; + T* R_warm = new T[n * n](); + RandLAPACK::CQRRT_linops warm_algo(false, tol, false); + warm_algo.nnz = sketch_nnz; + warm_algo.block_size = block_size; + warm_algo.call(V_app_op, R_warm, n, d_factor, warm_state); + delete[] R_warm; + } + std::cout << "done\n\n"; + + std::vector> all_results; + + // Per-iteration QR output; invariant size across all (alg, run) iters. + T* R = new T[n * n](); + + for (const auto& alg_name : selected_algs) { + std::cout << "\n=== Algorithm: " << alg_name << " (rspec) ===\n"; + + for (int64_t run_idx = 0; run_idx < num_runs; ++run_idx) { + rspec_result res{}; + res.m = m; + res.n = n; + res.run_idx = run_idx; + res.alg_name = alg_name; + res.qr_status = 0; + res.qr_time_us = 0; + res.peak_rss_kb = 0; + res.analytical_kb = -1; // -1 = no value; 0 would read as a real 0 MB bar + res.factor_time_us = factor_time_us; + res.rspec_total_us = 0; + res.orth_error = (T)-1.0; + + auto rspec_t0 = steady_clock::now(); + + std::fill(R, R + n * n, (T)0); + auto state = run_states[run_idx]; + + std::cout << "[Run " << run_idx << ", " << alg_name << "] QR ... " << std::flush; + RandLAPACK::PeakRSSTracker mem; mem.start(); + if (alg_name == "sCholQR3") { + RandLAPACK::sCholQR3_linops qr_algo(true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + qr_algo.block_size = block_size; + res.qr_status = qr_algo.call(V_app_op, R, n); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::scholqr3_linops_analytical_kb(m, n, block_size); + } + } else if (alg_name == "sCholQR3_basic") { + RandLAPACK::sCholQR3_linops_basic qr_algo(true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + res.qr_status = qr_algo.call(V_app_op, R, n); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::scholqr3_linops_basic_analytical_kb(m, n); + } + } else if (alg_name == "CholQR") { + RandLAPACK::CholQR_linops qr_algo(true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + qr_algo.block_size = block_size; + res.qr_status = qr_algo.call(V_app_op, R, n); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector (6 entries; writer pads) + res.analytical_kb = RandLAPACK::cholqr_linops_analytical_kb(m, n, block_size); + } + } else if (alg_name == "CholQR2") { + RandLAPACK::CholQR2_linops qr_algo(true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + qr_algo.block_size = block_size; + res.qr_status = qr_algo.call(V_app_op, R, n); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::cholqr2_linops_analytical_kb(m, n, block_size); + } + } else { + // CQRRT_linop. CQRRT_linop_bqrrp is not part of the dispatch. + RandLAPACK::CQRRT_linops qr_algo(true, tol); + qr_algo.max_retries = bench_chol_max_retries(); + qr_algo.nnz = sketch_nnz; + qr_algo.block_size = block_size; + qr_algo.precond_method = RandLAPACK::CQRRTLinopPrecond::TRSM_IDENTITY; + res.qr_status = qr_algo.call(V_app_op, R, n, d_factor, state); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown = qr_algo.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::cqrrt_linops_analytical_kb(m, n, d_factor, block_size); + } + } + + if (res.qr_status != 0) { + std::cerr << "\n [" << alg_name << "] Run " << run_idx + << ": QR returned status " << res.qr_status + << ". Skipping eigen post-processing.\n"; + res.qr_time_us = -1; + res.qr_breakdown.clear(); // writer pads to 18 zero columns regardless of size + res.rr_breakdown.clear(); + res.orth_error = std::numeric_limits::quiet_NaN(); + res.top_eigvals.assign(top_k, std::numeric_limits::quiet_NaN()); + res.top_residuals.assign(top_k, std::numeric_limits::quiet_NaN()); + auto rspec_t1 = steady_clock::now(); + res.rspec_total_us = duration_cast(rspec_t1 - rspec_t0).count(); + all_results.push_back(res); + continue; + } + std::cout << "done (" << res.qr_time_us << " us)\n"; + + // ---- Orthogonality loss of the Q-factor: ||Q^T Q - I||_F / sqrt(n), + // Q = V_app * R^{-1}, materialized explicitly (same path as irlsq). + // NOTE: this re-applies V_app to n columns (one extra full pass of the + // C^j chain); at FEM2 scale that is a meaningful cost. + steady_clock::time_point rr_t0, rr_t1; + long orth_us = 0, rr_build_us = 0, syevd_us = 0, resid_us = 0; + + std::cout << " orth loss ... " << std::flush; + rr_t0 = steady_clock::now(); + res.orth_error = compute_orth_error_explicit(V_app_op, R, m, n, block_size); + rr_t1 = steady_clock::now(); + orth_us = duration_cast(rr_t1 - rr_t0).count(); + std::cout << "done (" << std::scientific << std::setprecision(3) + << res.orth_error << ")\n"; + + // ---------------------------------------------------------------- + // Rayleigh-Ritz: T = R^{-T} V_app^T C V_app R^{-1} + // Materialize V_app^T C V_app column-block by column-block on identity. + // ---------------------------------------------------------------- + std::cout << " Building Rayleigh-Ritz matrix T (n=" << n << ") ... " << std::flush; + rr_t0 = steady_clock::now(); + int64_t b_rr = (block_size > 0) ? std::min(block_size, n) : std::min(64, n); + + T* T_mat = new T[(size_t)n * (size_t)n](); + T* eye_blk = new T[(size_t)n * (size_t)b_rr](); + T* Va_blk = new T[(size_t)m * (size_t)b_rr](); + T* CV_blk = new T[(size_t)m * (size_t)b_rr](); + T* T_blk = new T[(size_t)n * (size_t)b_rr](); + + for (int64_t j0 = 0; j0 < n; j0 += b_rr) { + int64_t bk = std::min(b_rr, n - j0); + + std::fill_n(eye_blk, (size_t)n * (size_t)b_rr, (T)0); + for (int64_t j = 0; j < bk; ++j) + eye_blk[(j0 + j) + j * n] = (T)1; + + // Va_blk = V_app * eye_blk (m x bk) + V_app_op(blas::Side::Left, blas::Layout::ColMajor, + blas::Op::NoTrans, blas::Op::NoTrans, + m, bk, n, (T)1.0, eye_blk, n, (T)0.0, Va_blk, m); + + // CV_blk = C * Va_blk (m x bk) + C_op(blas::Side::Left, blas::Layout::ColMajor, + blas::Op::NoTrans, blas::Op::NoTrans, + m, bk, m, (T)1.0, Va_blk, m, (T)0.0, CV_blk, m); + + // T_blk = V_app^T * CV_blk (n x bk) + V_app_op(blas::Side::Left, blas::Layout::ColMajor, + blas::Op::Trans, blas::Op::NoTrans, + n, bk, m, (T)1.0, CV_blk, m, (T)0.0, T_blk, n); + + lapack::lacpy(lapack::MatrixType::General, n, bk, T_blk, n, T_mat + j0 * n, n); + } + + delete[] eye_blk; + delete[] Va_blk; + delete[] CV_blk; + delete[] T_blk; + + // Apply R^{-T} on the left: T := R^{-T} * T + blas::trsm(blas::Layout::ColMajor, blas::Side::Left, blas::Uplo::Upper, + blas::Op::Trans, blas::Diag::NonUnit, + n, n, (T)1.0, R, n, T_mat, n); + // Apply R^{-1} on the right: T := T * R^{-1} + blas::trsm(blas::Layout::ColMajor, blas::Side::Right, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, + n, n, (T)1.0, R, n, T_mat, n); + + // Symmetrize against rounding drift before the (upper-triangle) syevd. + RandBLAS::symmetrize(blas::Layout::ColMajor, blas::Uplo::Upper, n, T_mat, n); + rr_t1 = steady_clock::now(); + rr_build_us = duration_cast(rr_t1 - rr_t0).count(); + std::cout << "done\n"; + + // Eigendecomposition: T = U diag(lambda) U^T, U overwrites T_mat (columns = eigvecs). + std::cout << " syevd ... " << std::flush; + rr_t0 = steady_clock::now(); + T* eigvals = new T[(size_t)n](); + int64_t syevd_info = lapack::syevd(lapack::Job::Vec, blas::Uplo::Upper, + n, T_mat, n, eigvals); + if (syevd_info != 0) { + std::cerr << "Warning: syevd returned " << syevd_info << " for run " << run_idx << "\n"; + } + // syevd returns eigvals in ascending order; collect top-k by absolute magnitude (largest |lambda|). + std::cout << "done\n"; + + // Sort eigenvalues by descending |lambda|; build permutation. + int64_t* perm = new int64_t[(size_t)n]; + for (int64_t i = 0; i < n; ++i) perm[i] = i; + std::sort(perm, perm + n, [&](int64_t a, int64_t b_){ + return std::abs(eigvals[a]) > std::abs(eigvals[b_]); + }); + + res.top_eigvals.resize(top_k); + for (int i = 0; i < top_k; ++i) res.top_eigvals[i] = eigvals[perm[i]]; + rr_t1 = steady_clock::now(); + syevd_us = duration_cast(rr_t1 - rr_t0).count(); + + // ---------------------------------------------------------------- + // Ritz residual norms for the top-k pairs (collaborator spec): + // resid_i = ||C y_i - lambda_i y_i|| / (|lambda_max| * ||y_i||) + // where y_i = Q u_i = V_app R^{-1} u_i is the Ritz vector and u_i is an + // eigenvector of the small RR matrix T = Q^T C Q (Q = V_app R^{-1}, with + // Q^T Q = I via the Q-less QR). This is the ordinary (non-generalized) eigen- + // residual of the symmetric operator C, which is exactly what Rayleigh- + // Ritz on range(V_app) approximates. |lambda_max| is the dominant Ritz + // value (largest |lambda|), used as the relative scale. K and M are no + // longer needed here, only C and the Ritz vectors. + // ---------------------------------------------------------------- + std::cout << " Ritz residuals ... " << std::flush; + rr_t0 = steady_clock::now(); + + // u_blk = R^{-1} * U_topk (n x top_k); columns = top-k eigenvectors of T. + T* u_blk = new T[(size_t)n * (size_t)top_k](); + for (int i = 0; i < top_k; ++i) + for (int64_t r = 0; r < n; ++r) + u_blk[r + i * n] = T_mat[r + perm[i] * n]; + blas::trsm(blas::Layout::ColMajor, blas::Side::Left, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, + n, top_k, (T)1.0, R, n, u_blk, n); + + // y_blk = V_app * (R^{-1} U_topk) = Q U_topk (m x top_k): the Ritz vectors. + T* y_blk = new T[(size_t)m * (size_t)top_k](); + V_app_op(blas::Side::Left, blas::Layout::ColMajor, + blas::Op::NoTrans, blas::Op::NoTrans, + m, top_k, n, (T)1.0, u_blk, n, (T)0.0, y_blk, m); + + // Cy_blk = C * y_blk (m x top_k). + T* Cy_blk = new T[(size_t)m * (size_t)top_k](); + C_op(blas::Side::Left, blas::Layout::ColMajor, + blas::Op::NoTrans, blas::Op::NoTrans, + m, top_k, m, (T)1.0, y_blk, m, (T)0.0, Cy_blk, m); + + // |lambda_max| = dominant Ritz value (top_eigvals sorted by descending |lambda|). + T lam_max = (top_k > 0) ? std::abs(res.top_eigvals[0]) : (T)0; + + res.top_residuals.resize(top_k); + for (int i = 0; i < top_k; ++i) { + T lam = res.top_eigvals[i]; + T num_sq = 0, y_sq = 0; + for (int64_t r = 0; r < m; ++r) { + T d = Cy_blk[r + i * m] - lam * y_blk[r + i * m]; + num_sq += d * d; + y_sq += y_blk[r + i * m] * y_blk[r + i * m]; + } + T denom = lam_max * std::sqrt(y_sq); + res.top_residuals[i] = (denom > 0) ? std::sqrt(num_sq) / denom + : std::numeric_limits::quiet_NaN(); + } + rr_t1 = steady_clock::now(); + resid_us = duration_cast(rr_t1 - rr_t0).count(); + + delete[] u_blk; + delete[] y_blk; + delete[] Cy_blk; + delete[] eigvals; + delete[] perm; + delete[] T_mat; + std::cout << "done\n"; + + auto rspec_t1 = steady_clock::now(); + // Exclude the orth-loss diagnostic from the headline total: it is a + // pure verification quantity costing a full extra n-column pass of + // the C^j chain and would otherwise contaminate rspec_total_us. It + // stays visible as rr_breakdown[0]. + res.rspec_total_us = duration_cast(rspec_t1 - rspec_t0).count() - orth_us; + res.rr_breakdown = {orth_us, rr_build_us, syevd_us, resid_us}; + + std::cout << " Top eigvals: "; + for (int i = 0; i < std::min(5, top_k); ++i) + std::cout << res.top_eigvals[i] << " "; + std::cout << "\n"; + std::cout << " Top residuals: "; + for (int i = 0; i < std::min(5, top_k); ++i) + std::cout << res.top_residuals[i] << " "; + std::cout << "\n"; + + all_results.push_back(res); + } + } + + // CSV output + std::string time_buf = make_run_timestamp(); + std::string results_file = output_dir + "/" + time_buf + "_rspec_results.csv"; + write_rspec_csv(results_file, all_results, m, n, num_runs, + K_file, M_file, V_file, omega, power_j, + sketch_nnz, block_size, method_mask, top_k); + std::cout << "\nRSPEC results written to " << results_file << "\n"; + + std::string breakdown_file = output_dir + "/" + time_buf + "_rspec_breakdown.csv"; + write_rspec_breakdown(breakdown_file, all_results); + std::cout << "RSPEC breakdown written to " << breakdown_file << "\n"; + + delete[] R; + return 0; +} + +// ============================================================================ +// CSV writer: IR-LSQ regularized (irlsq_reg): base columns + regularization / +// mixed-precision metadata (kappa_target, kappa_measured, mu, precond/solve prec) +// ============================================================================ + +template +static void write_irlsq_reg_results( + const std::string& filename, + const std::vector>& results, + int64_t m, int64_t n, int64_t nnz_or_zero, const std::string& input_label, + double d_factor, int64_t sketch_nnz, int64_t block_size, int64_t method_mask, + double kappa_target, double mu, + const std::string& precond_prec, const std::string& solve_prec, + int64_t num_runs, long chol_time_us, double noise_level) +{ + std::ofstream out(filename); + // irlsq_reg is FEM-only (main() rejects sparse input for this mode), unlike + // the plain irlsq writer above, which serves both; the title says so. + out << "# FEM IR-LSQ (regularized augmented operator) Benchmark results\n" + << "# Date: " << make_run_timestamp() << "\n" + << "# argv=" << g_argv_line << "\n" + << "# input=" << input_label << "\n" + << "# M=" << m << " N=" << n << " nnz=" << nnz_or_zero << "\n" + << "# noise_level=" << noise_level << "\n" + << "# chol_time_us=" << chol_time_us << "\n" + << "# d_factor=" << d_factor << " sketch_nnz=" << sketch_nnz + << " block_size=" << block_size << "\n" + << "# method_mask=" << method_mask << "\n" + << "# num_runs=" << num_runs << "\n" + << "# kappa_target=" << kappa_target << " mu=" << mu << "\n" + << "# precond_prec=" << precond_prec << " solve_prec=" << solve_prec << "\n" + << "# blendenpik=warm+cold (IR methods always cold x0);" + << " refine rows = init_only x0 + shared engine\n" + << "# A_hat = [A; mu*I]; R = chol(A^T A + mu^2 I) built in precond_prec,\n" + << "# used as right preconditioner for IterRefineLSQ run in solve_prec.\n" +#ifdef _OPENMP + << "# OpenMP threads: " << omp_get_max_threads() << "\n" +#else + << "# OpenMP threads: 1\n" +#endif + ; + write_env_provenance(out); + out << "algorithm,run,m,n,qr_status,qr_time_us,peak_rss_kb,analytical_kb," + "orth_error,ir_total_us,ir_outer_iters,ir_inner_iters_total," + "ls_residual_norm,ls_solution_error,kappa_target,kappa_measured,mu,precond_prec,solve_prec,chol_retries," + "ir_inner_capped,ir_inner_relres,ir_inner_best_relres,ir_inner_best_iter,cond_precond,ir_setup_us," + "lsqr_iters,engine_status,stop_reason,x0_relres,chol_shift_abs,chol_shift_rel\n"; + // Sentinel note: chol_retries and chol_shift_abs/chol_shift_rel use -1 for + // "no Cholesky in this row" (Blendenpik family, unpreconditioned) or "QR + // failed before a retry/shift record existed"; 0 still means "Cholesky + // ran unshifted" (chol_retries) or "ran, no shift applied" (the shifts). + for (const auto& r : results) { + out << r.alg_name << "," << r.run_idx << "," << r.m << "," << r.n << "," + << r.qr_status << "," << r.qr_time_us << "," << r.peak_rss_kb << "," << r.analytical_kb << "," + << std::scientific << std::setprecision(6) << r.orth_error << "," + << r.ir_total_us << "," << r.ir_outer_iters << "," << r.ir_inner_iters_total << "," + << std::scientific << std::setprecision(6) << r.ls_residual_norm << "," + << std::scientific << std::setprecision(6) << r.ls_solution_error << "," + << std::scientific << std::setprecision(6) << kappa_target << "," + << std::scientific << std::setprecision(6) << r.kappa_measured << "," + << std::scientific << std::setprecision(6) << mu << "," + << precond_prec << "," << solve_prec << "," << r.chol_retries << "," + << r.ir_inner_capped << "," + << std::scientific << std::setprecision(6) << r.ir_inner_relres << "," + << std::scientific << std::setprecision(6) << r.ir_inner_best_relres << "," + << r.ir_inner_best_iter << "," + << std::scientific << std::setprecision(6) << r.cond_precond << "," + << r.ir_setup_us << "," + << r.lsqr_iters << "," << r.engine_status << "," << r.stop_reason << "," + << std::scientific << std::setprecision(6) << r.x0_relres << "," + << std::scientific << std::setprecision(6) << r.chol_shift_abs << "," + << std::scientific << std::setprecision(6) << r.chol_shift_rel + << "\n"; + } +} + +// kappa(A) estimate from the regularized R diagonal: max|R_ii| / min|R_ii|. +template +static double kappa_from_R_diag(const P* R, int64_t n) { + double mx = 0.0, mn = std::numeric_limits::infinity(); + for (int64_t i = 0; i < n; ++i) { + double v = std::abs((double)R[i + i * n]); + if (v > mx) mx = v; + if (v > 0 && v < mn) mn = v; + } + return (mn > 0 && std::isfinite(mn)) ? mx / mn : -1.0; +} + +// ============================================================================ +// irlsq_reg runner: regularized augmented-operator preconditioner with +// independent preconditioner (P_precond) and solve (T_solve) precisions. +// +// Builds two FEM operator chains J = L^{-1} K (V D) from the same kappa-scaled +// matrices: one in P_precond (for Q-less QR of A_hat = [A; mu*I]) and one in +// T_solve (for IterRefineLSQ on the base A). For each variant: QR in P_precond +// -> R (= chol(A^T A + mu^2 I)) -> cast to T_solve -> solve. R is never stored +// for all variants at once (n^2 is huge at FEM2 scale), so QR and solve are +// interleaved and both chains coexist. +// ============================================================================ + +template +static int run_irlsq_reg( + const std::string& K_file, const std::string& M_file, const std::string& V_file, + const std::string& output_dir, int64_t num_runs, + double d_factor, int64_t sketch_nnz, int64_t block_size, + bool compute_cond, + int64_t method_mask, double kappa_target, double mu_factor, double noise_level, + const std::string& precond_prec_str, const std::string& solve_prec_str) +{ + namespace rl = RandLAPACK::linops; + + // Shared decode (see decode_method_mask and the file-header mask docs). + std::vector selected_algs = decode_method_mask(method_mask, /*with_blendenpik=*/true); + if (selected_algs.empty()) { + std::cerr << "Error: method_mask selects no algorithms (got " << method_mask << ").\n"; + return 1; + } + + // ---- Load double master CSRs ---- + int64_t m_K, n_K, nnz_K, m_M, n_M, nnz_M, m_V, n_V, nnz_V; + auto K_master = load_csr_verbose("K (stiffness)", K_file, m_K, n_K, nnz_K); + auto M_master = load_csr_verbose("M (mass)", M_file, m_M, n_M, nnz_M); + auto V_master = load_csr_verbose("V (prolongation)", V_file, m_V, n_V, nnz_V); + if (m_K != n_K) { std::cerr << "Error: K must be square.\n"; return 1; } + if (m_M != m_K || n_M != m_K) { std::cerr << "Error: M size must match K.\n"; return 1; } + if (m_V != m_K) { std::cerr << "Error: V rows must match K size.\n"; return 1; } + if (m_V < n_V) { std::cerr << "Error: need tall V (m_fine >= n_coarse).\n"; return 1; } + int64_t m = m_V, n = n_V; + + // ---- Inject conditioning: scale V columns by the geometric diagonal ---- + auto d_scale = geometric_colscale(n_V, kappa_target); + scale_csr_columns(V_master, d_scale); + std::cout << "Column-scaled V to target kappa=" << kappa_target + << " (spread " << d_scale.front() << " .. " << d_scale.back() << ")\n"; + + // ---- Build SOLVE chain (precision T_solve), cast down from double master ---- + auto K_Ts = csr_cast(K_master); + auto V_Ts = csr_cast(V_master); + auto M_Ts = csr_cast(M_master); + rl::SparseLinOp> K_op_Ts(m_K, m_K, K_Ts); + rl::SparseLinOp> V_op_Ts(m_V, n_V, V_Ts); + std::cout << "Factorizing M = L L^T (solve precision)... " << std::flush; + RandLAPACK_extras::linops::CholSolverLinOp L_inv_Ts(M_Ts, /*half_solve=*/true); + auto chol_ts_t0 = steady_clock::now(); + L_inv_Ts.factorize(); + auto chol_ts_t1 = steady_clock::now(); + std::cout << "done\n"; + rl::CompositeOperator KV_Ts(m, n, K_op_Ts, V_op_Ts); KV_Ts.block_size = block_size; + rl::CompositeOperator J_Ts(m, n, L_inv_Ts, KV_Ts); J_Ts.block_size = block_size; + + // Consistent RHS: x_true ~ U(-1,1)^n, b = A x_true (+ noise_level relative + // Gaussian noise). Consistency makes the residual metric a true backward error + // ~u (kappa-robust: the ||A|| ||x|| factor cancels), and x_true gives a ground + // -truth forward-error metric ||x - x_true|| / ||x_true|| ~ u*kappa that exposes + // the precision x kappa interaction. Use noise_level = 0 to see the solver's + // u-level backward error directly. (Same construction as sparse mode.) + std::vector x_true(n, (T_solve)0); + { std::mt19937 rng_x(42); std::uniform_real_distribution U(-1.0, 1.0); + for (auto& v : x_true) v = (T_solve)U(rng_x); } + std::vector b(m, (T_solve)0); + J_Ts(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T_solve)1.0, x_true.data(), n, (T_solve)0.0, b.data(), m); + if (noise_level > 0) { + T_solve b_clean_norm = blas::nrm2(m, b.data(), 1); + std::vector noise(m, (T_solve)0); + std::mt19937 rng_n(13); std::normal_distribution N01(0, 1); + for (auto& v : noise) v = (T_solve)N01(rng_n); + T_solve raw = blas::nrm2(m, noise.data(), 1); + T_solve scale = (raw > 0) ? (T_solve)(noise_level) * b_clean_norm / raw : (T_solve)0; + for (int64_t i = 0; i < m; ++i) b[i] += scale * noise[i]; + } + const T_solve x_true_norm = blas::nrm2(n, x_true.data(), 1); + std::cout << "Consistent RHS b = A x_true" + << (noise_level > 0 ? " + noise" : "") + << " (||x_true||=" << x_true_norm << ", ||b||=" << blas::nrm2(m, b.data(), 1) << ")\n"; + + // ---- Build PRECOND chain (precision P_precond) + augmented operator ---- + auto K_Pp = csr_cast(K_master); + auto V_Pp = csr_cast(V_master); + auto M_Pp = csr_cast(M_master); + rl::SparseLinOp> K_op_Pp(m_K, m_K, K_Pp); + rl::SparseLinOp> V_op_Pp(m_V, n_V, V_Pp); + std::cout << "Factorizing M = L L^T (precond precision)... " << std::flush; + RandLAPACK_extras::linops::CholSolverLinOp L_inv_Pp(M_Pp, /*half_solve=*/true); + auto chol_pp_t0 = steady_clock::now(); + L_inv_Pp.factorize(); + auto chol_pp_t1 = steady_clock::now(); + std::cout << "done\n"; + rl::CompositeOperator KV_Pp(m, n, K_op_Pp, V_op_Pp); KV_Pp.block_size = block_size; + rl::CompositeOperator J_Pp(m, n, L_inv_Pp, KV_Pp); J_Pp.block_size = block_size; + + // Shared per-cell cost: both M-factorizations (solve + precond + // precision), measured once, identical for every row in this run. + const long chol_time_us = duration_cast(chol_ts_t1 - chol_ts_t0).count() + + duration_cast(chol_pp_t1 - chol_pp_t0).count(); + + // ||A||_2 and ||b|| for the Higham backward-error metric. + T_solve A_2norm = estimate_op_2norm(J_Ts, m, n, 10); + T_solve b_norm = blas::nrm2(m, b.data(), 1); + std::cout << "||A||_2 ~ " << A_2norm << ", ||b|| = " << b_norm << "\n"; + + // Regularization per the collaborator's spec: mu = mu_factor * u(precond), + // with mu_factor = 10 giving mu = 10u (u = unit roundoff of the precond + // precision). NO ||A|| or size scaling: the augmented operator is exactly + // A_hat = [A; mu*I], Q-less CholeskyQR of which gives R = chol(A^T A + mu^2 I), + // used as a right preconditioner for the LS problem in A. + const P_precond mu_P = (P_precond)(mu_factor * (double)unit_roundoff()); + rl::ScaledIdentityOp reg_op(n, mu_P); + rl::VStackOp> A_hat_Pp(J_Pp, reg_op); + A_hat_Pp.block_size = block_size; // caps the blocked-sketch slice width (CQRRT) + std::cout << "Augmented operator A_hat = [J; mu*I], mu=" << (double)mu_P + << " (= " << mu_factor << " * u(" << precond_prec_str << "))\n\n"; + + const P_precond tol_P = std::pow(std::numeric_limits::epsilon(), (P_precond)0.85); + const T_solve tol_T = std::pow(std::numeric_limits::epsilon(), (T_solve)0.85); + + // Per-run RNG states (CQRRT only). + RandBLAS::RNGState main_state(123); + std::vector> run_states(num_runs); + for (int64_t r = 0; r < num_runs; ++r) { run_states[r] = main_state; if (r > 0) run_states[r].key.incr(r); } + + // Warmup the precond-chain CQRRT on A_hat (warms the L^{-1} K V chain, the + // augmented Gram, and the blocked sketch overload), then the SOLVE chain: + // the timed IR-LSQ runs LSQR on J_Ts with a TRSM preconditioner, and its + // thread pools / first-touch pages otherwise land inside the FIRST + // method's timed solve, which at 4-7 inner iterations is the same magnitude + // as the whole solve. A few untimed LSQR iterations + // on J_Ts (with the warmup R when usable) close that gap. This is a CPU + // warmup, distinct from Blendenpik's x0 warm start. + std::cout << "Running warmup... " << std::flush; + { auto ws = run_states[0]; P_precond* Rw = new P_precond[n * n](); + RandLAPACK::CQRRT_linops warm(false, tol_P); + warm.nnz = sketch_nnz; warm.block_size = block_size; + int warm_status = warm.call(A_hat_Pp, Rw, n, (P_precond)d_factor, ws); + T_solve* Rw_T = new T_solve[n * n]; + if (warm_status == 0) + for (int64_t i = 0; i < n * n; ++i) Rw_T[i] = (T_solve)Rw[i]; + T_solve* x_wu = new T_solve[n](); + int it_wu = 0; long lt_wu[4] = {0}; + RandLAPACK::lsqr(J_Ts, m, n, + (warm_status == 0) ? Rw_T : nullptr, (warm_status == 0) ? n : (int64_t)0, + b.data(), x_wu, tol_T, tol_T, 5, it_wu, lt_wu); + delete[] Rw; delete[] Rw_T; delete[] x_wu; } + std::cout << "done\n"; + + std::vector> all_results; + + // Both n^2 buffers must be pre-touched before the measured loop starts. + // Value-init ("()") already zero-fills R_P, which pages it in + // here; R_T gets an explicit std::fill for the same effect (kept separate + // from allocation so the pre-touch is self-documenting rather than relying + // on new[]() semantics matching by accident on the next edit). + P_precond* R_P = new P_precond[n * n](); + T_solve* R_T = new T_solve[n * n]; + std::fill(R_T, R_T + n * n, (T_solve)0); + T_solve* x_ls = new T_solve[n]; + + for (const auto& alg_name : selected_algs) { + std::cout << "\n=== Algorithm: " << alg_name << " (irlsq_reg) ===\n"; + for (int64_t run_idx = 0; run_idx < num_runs; ++run_idx) { + bench_result res{}; + res.m = m; res.n = n; res.run_idx = run_idx; res.alg_name = alg_name; + res.qr_status = 0; res.qr_time_us = 0; res.orth_error = (T_solve)-1; + res.ls_residual_norm = (T_solve)-1; res.ls_solution_error = (T_solve)-1; + res.kappa_measured = (T_solve)-1; + res.chol_time_us = chol_time_us; // shared per-cell cost + + std::fill(R_P, R_P + n * n, (P_precond)0); + auto state = run_states[run_idx]; + const bool is_bp = (alg_name.rfind("Blendenpik", 0) == 0); + + std::cout << "[Run " << run_idx << ", " << alg_name << "] QR(" << precond_prec_str + << ") ... " << std::flush; + RandLAPACK::PeakRSSTracker mem; mem.start(); + if (is_bp) { + // Shared Blendenpik-family dispatch, in SOLVE precision on the BASE + // operator J_Ts (no mu, no augmented A_hat); fills every accounting + // field itself and writes the sketch R factor into R_T directly. + T_solve outer_tol_eff = (g_ir_outer_tol >= 0) ? (T_solve)g_ir_outer_tol + : (T_solve)10 * std::numeric_limits::epsilon(); + run_blendenpik_family(J_Ts, b.data(), m, x_ls, n, R_T, + (T_solve)d_factor, sketch_nnz, state, alg_name, tol_T, outer_tol_eff, + mem, res); + } else if (alg_name == "sCholQR3") { + RandLAPACK::sCholQR3_linops qr(true, tol_P); qr.block_size = block_size; + qr.max_retries = bench_chol_max_retries(); + res.qr_status = qr.call(A_hat_Pp, R_P, n); res.chol_retries = qr.n_chol_retries; + record_chol_shift(res, qr.chol_applied_shifts, qr.chol_gram_traces); + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown = qr.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::scholqr3_linops_analytical_kb(A_hat_Pp.n_rows, n, block_size); } + } else if (alg_name == "sCholQR3_basic") { + RandLAPACK::sCholQR3_linops_basic qr(true, tol_P); + qr.max_retries = bench_chol_max_retries(); + res.qr_status = qr.call(A_hat_Pp, R_P, n); res.chol_retries = qr.n_chol_retries; + record_chol_shift(res, qr.chol_applied_shifts, qr.chol_gram_traces); + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown = qr.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::scholqr3_linops_basic_analytical_kb(A_hat_Pp.n_rows, n); } + } else if (alg_name == "CholQR") { + RandLAPACK::CholQR_linops qr(true, tol_P); qr.block_size = block_size; + qr.max_retries = bench_chol_max_retries(); + res.qr_status = qr.call(A_hat_Pp, R_P, n); res.chol_retries = qr.n_chol_retries; + record_chol_shift(res, qr.chol_applied_shifts, qr.chol_gram_traces); + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown = qr.times; // whole vector (6 entries; writer pads) + res.analytical_kb = RandLAPACK::cholqr_linops_analytical_kb(A_hat_Pp.n_rows, n, block_size); } + } else if (alg_name == "CholQR2") { + RandLAPACK::CholQR2_linops qr(true, tol_P); qr.block_size = block_size; + qr.max_retries = bench_chol_max_retries(); + res.qr_status = qr.call(A_hat_Pp, R_P, n); res.chol_retries = qr.n_chol_retries; + record_chol_shift(res, qr.chol_applied_shifts, qr.chol_gram_traces); + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown = qr.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::cholqr2_linops_analytical_kb(A_hat_Pp.n_rows, n, block_size); } + } else { + // CQRRT: sketch + Gram the augmented A_hat (via VStack's blocked sketch + // overload), uniformly with the other 4 methods. R = chol(A^T A + mu^2 I). + RandLAPACK::CQRRT_linops qr(true, tol_P); + qr.max_retries = bench_chol_max_retries(); + qr.nnz = sketch_nnz; qr.block_size = block_size; + qr.precond_method = RandLAPACK::CQRRTLinopPrecond::TRSM_IDENTITY; + res.qr_status = qr.call(A_hat_Pp, R_P, n, (P_precond)d_factor, state); res.chol_retries = qr.n_chol_retries; + record_chol_shift(res, qr.chol_applied_shifts, qr.chol_gram_traces); + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown = qr.times; // whole vector, not truncated to a fixed slot count + res.analytical_kb = RandLAPACK::cqrrt_linops_analytical_kb(A_hat_Pp.n_rows, n, (P_precond)d_factor, block_size); } + } + + if (res.qr_status != 0) { + std::cerr << "\n [" << alg_name << "] Run " << run_idx + << ": QR returned status " << res.qr_status << ". Skipping solve.\n"; + res.qr_time_us = -1; res.ir_total_us = -1; res.qr_breakdown.clear(); res.analytical_kb = -1; // -1 = no value; 0 would read as a real 0 MB bar + if (!is_bp) res.peak_rss_kb = mem.stop(); // bp rows already stopped inside the helper + all_results.push_back(res); + continue; + } + res.kappa_measured = is_bp ? (T_solve)kappa_from_R_diag(R_T, n) + : (T_solve)kappa_from_R_diag(R_P, n); + std::cout << "done (" << res.qr_time_us << " us, kappa~" + << std::scientific << std::setprecision(2) << (double)res.kappa_measured << ")"; + + // Cast R to solve precision (Blendenpik already produced R_T directly). + if (!is_bp) for (int64_t i = 0; i < n * n; ++i) R_T[i] = (T_solve)R_P[i]; + + // NOTE ordering: the orthogonality diagnostic is computed + // AFTER the solve, not here. It materializes Q = A R^{-1} (m x n, about + // 88 GB on the large cell), so leaving it inside the peak-RSS window + // would swamp the measurement. Moving it below lets every row family + // use the SAME window, build + solve with diagnostics excluded, which + // is what the Toeplitz benchmark already did and what makes the + // peak-vs-predicted panel comparable across the two figures. + + // Solve in solve precision. Blendenpik-family rows already solved and + // recorded everything inside run_blendenpik_family; everyone else runs + // IR-LSQ with R as the right preconditioner. + if (is_bp) { + std::cout << ". solve(" << solve_prec_str << ") recorded ... " << std::flush; + } else { + std::cout << ". IR-LSQ(" << solve_prec_str << ") ... " << std::flush; + auto ls_t0 = steady_clock::now(); + // Always cold: warm x0 is Blendenpik-only (see the CLI comment + // block). ir_setup_us stays in the CSV schema and is always 0 + // for IR methods. + std::fill(x_ls, x_ls + n, (T_solve)0.0); + RandLAPACK::IterRefineLSQ ir( + (g_ir_inner_tol > 0) ? (T_solve)g_ir_inner_tol : tol_T, + (g_ir_max_inner > 0) ? g_ir_max_inner : 200, + g_ir_n_steps, true, false); + ir.round_drop = (T_solve)g_ir_round_drop; + ir.outer_tol = (g_ir_outer_tol >= 0) ? (T_solve)g_ir_outer_tol + : (T_solve)10 * std::numeric_limits::epsilon(); + ir.outer_stag_window = ir_outer_stag_window(); + int ir_status = ir.call(J_Ts, R_T, n, b.data(), m, x_ls, n); + auto ls_t1 = steady_clock::now(); + if (ir_status != 0) std::cerr << "Warning: IterRefineLSQ status " << ir_status << "\n"; + res.ir_total_us = duration_cast(ls_t1 - ls_t0).count(); + res.ir_outer_iters = ir.outer_iters_done; + res.ir_inner_iters_total = 0; + for (int v : ir.inner_iters_per_step) res.ir_inner_iters_total += v; + record_inner_cg_diagnosis(ir, res); + record_ir_outputs(ir, res); + if (!ir.times.empty()) res.ir_breakdown = ir.times; + res.peak_rss_kb = mem.stop(); // build + solve, diagnostics excluded + } + + // Orthogonality loss of Q = A R^{-1} (base A in solve precision). + // Outside the RSS window on purpose; see the ordering note above. + // compute_cond gates cond_precond identically in both irlsq and + // irlsq_reg. + res.orth_error = compute_orth_error_explicit(J_Ts, R_T, m, n, block_size, + compute_cond ? &res.cond_precond : nullptr); + + // Higham normwise backward error ||Ax-b|| / (||A||_2 ||x|| + ||b||). + std::vector Ax(m, (T_solve)0); + J_Ts(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T_solve)1.0, x_ls, n, (T_solve)0.0, Ax.data(), m); + T_solve resid_sq = 0; + for (int64_t i = 0; i < m; ++i) { T_solve dd = Ax[i] - b[i]; resid_sq += dd * dd; } + T_solve x_norm = blas::nrm2(n, x_ls, 1); + T_solve denom = A_2norm * x_norm + b_norm; + res.ls_residual_norm = (denom > 0) ? std::sqrt(resid_sq) / denom : (T_solve)-1; + + // Forward error vs ground truth: ||x - x_true|| / ||x_true|| ~ u*kappa. + T_solve err_sq = 0; + for (int64_t i = 0; i < n; ++i) { T_solve dd = x_ls[i] - x_true[i]; err_sq += dd * dd; } + res.ls_solution_error = (x_true_norm > 0) ? std::sqrt(err_sq) / x_true_norm : (T_solve)-1; + + std::cout << "done (" << res.ir_total_us << " us, bwd_err=" + << std::scientific << std::setprecision(3) << (double)res.ls_residual_norm + << ", fwd_err=" << (double)res.ls_solution_error << ")\n"; + + all_results.push_back(res); + } + } + delete[] R_P; delete[] R_T; delete[] x_ls; + + std::string time_buf = make_run_timestamp(); + std::string results_file = output_dir + "/" + time_buf + "_irlsq_reg_results.csv"; + std::string breakdown_file = output_dir + "/" + time_buf + "_irlsq_reg_breakdown.csv"; + std::string rounds_file = output_dir + "/" + time_buf + "_irlsq_reg_rounds.csv"; + write_irlsq_reg_results(results_file, all_results, m, n, nnz_K, + "L^{-1} K (V D) (M=" + M_file + ")", d_factor, sketch_nnz, block_size, method_mask, + kappa_target, (double)mu_P, precond_prec_str, solve_prec_str, + num_runs, chol_time_us, noise_level); + std::cout << "\nIR-LSQ-reg results written to " << results_file << "\n"; + write_irlsq_breakdown(breakdown_file, all_results, + "IR-LSQ (regularized augmented operator)"); + std::cout << "IR-LSQ-reg breakdown written to " << breakdown_file << "\n"; + write_rounds_csv(rounds_file, all_results); + std::cout << "IR-LSQ-reg per-round records written to " << rounds_file << "\n"; + return 0; +} + +// Dispatch run_irlsq_reg on the runtime precond-precision string (solve precision = T). +template +static int dispatch_irlsq_reg( + const std::string& precond_prec, + const std::string& K_file, const std::string& M_file, const std::string& V_file, + const std::string& output_dir, int64_t num_runs, + double d_factor, int64_t sketch_nnz, int64_t block_size, + bool compute_cond, + int64_t method_mask, double kappa_target, double mu_factor, double noise_level, + const std::string& solve_prec_str) +{ + if (precond_prec == "double") { + return run_irlsq_reg( + K_file, M_file, V_file, output_dir, num_runs, d_factor, sketch_nnz, block_size, + compute_cond, method_mask, kappa_target, mu_factor, noise_level, "double", solve_prec_str); + } else if (precond_prec == "single" || precond_prec == "float") { + return run_irlsq_reg( + K_file, M_file, V_file, output_dir, num_runs, d_factor, sketch_nnz, block_size, + compute_cond, method_mask, kappa_target, mu_factor, noise_level, "single", solve_prec_str); + } + std::cerr << "Error: precond_prec must be 'single' or 'double'; got '" << precond_prec << "'\n"; + return 1; +} + +// ============================================================================ +// Main dispatcher +// ============================================================================ + +template +int run_benchmark(int argc, char* argv[]) { + g_argv_line = quote_join_argv(argc, argv); + // ... → argc >= 5 to reach + if (argc < 8) { + std::cerr << "Usage: " << argv[0] + << " \n" + << " sparse mode: 'sparse' " + << " [sketch_nnz] [block_size] [compute_cond] [method_mask] [noise_level]\n" + << " FEM mode: " + << " [sketch_nnz] [block_size] [compute_cond] [method_mask] [noise_level]" + << " [omega] [power_j]\n" + << " mode = irlsq | rspec | irlsq_reg (rspec/irlsq_reg are FEM-only)\n"; + return 1; + } + + std::string output_dir = argv[2]; + int64_t num_runs = std::stoll(argv[3]); + std::string mode = argv[4]; + if (mode != "irlsq" && mode != "rspec" && mode != "irlsq_reg") { + std::cerr << "Error: must be one of {irlsq, rspec, irlsq_reg}; got '" << mode << "'\n"; + return 1; + } + + std::string arg5 = argv[5]; + bool sparse_mode = (arg5 == "sparse"); + + std::string K_file, M_file, V_file, A_file; + T d_factor; + int dfactor_idx; + + if (sparse_mode) { + if (argc < 8) { + std::cerr << "Error: sparse mode needs \n"; + return 1; + } + A_file = argv[6]; + d_factor = std::stod(argv[7]); + dfactor_idx = 7; + } else { + if (argc < 9) { + std::cerr << "Error: FEM mode needs \n"; + return 1; + } + K_file = arg5; + M_file = argv[6]; + V_file = argv[7]; + d_factor = std::stod(argv[8]); + dfactor_idx = 8; + } + + auto opt_long = [&](int rel, int64_t def) { + int idx = dfactor_idx + rel; + return (argc > idx) ? std::stoll(argv[idx]) : def; + }; + auto opt_double = [&](int rel, double def) { + int idx = dfactor_idx + rel; + return (argc > idx) ? std::stod(argv[idx]) : def; + }; + int64_t sketch_nnz = opt_long(1, 4); + // Default block_size=256 matches the paper's b=256 (blocked Gram); pass 0 + // explicitly for unblocked. + int64_t block_size = opt_long(2, 256); + bool compute_cond = (opt_long(3, 0) != 0); + int64_t method_mask = opt_long(4, 31); // see the file-header mask docs (bit 32 + // published Blendenpik, bit 64 refined; campaign mask 127) + const bool noise_level_explicit = (argc > dfactor_idx + 5); + T noise_level = (T)opt_double(5, 0.05); + double omega = opt_double(6, 0.0); + int64_t power_j = opt_long(7, 1); + // irlsq_reg-only knobs (positions after rspec's omega/power_j): + double kappa_target = opt_double(8, 1.0); // V column-scaling spread (1 = native) + double mu_factor = opt_double(9, 10.0); // mu = mu_factor * u(precond_prec) + std::string precond_prec = (argc > dfactor_idx + 10) ? std::string(argv[dfactor_idx + 10]) : "single"; + // Inner-CG controls (both optional and backward compatible). + // ir_max_inner <= 0 keeps the 200 default; ir_inner_tol < 0 keeps eps^0.85. + g_ir_max_inner = (int)opt_long(11, 200); + g_ir_inner_tol = opt_double(12, -1.0); + // Per-round CG residual drop (restart pacing). Slot 13 + // previously carried ir_inner_restarts, which the paced scheme obsoletes + // (every round IS a true-residual restart). Reject values >= 1 so a stale + // script passing the old integer restart count fails loudly rather than + // silently running near-empty rounds. 0 restores legacy fixed-tol rounds. + g_ir_round_drop = opt_double(13, 1e-4); + if (g_ir_round_drop < 0.0 || g_ir_round_drop >= 1.0) { + std::cerr << "Error: slot 13 is [ir_round_drop] (was " + "[ir_inner_restarts]); it must lie in [0, 1). Regenerate " + "the job scripts.\n"; + return 1; + } + // Outer-round cap. Campaign-canonical 50 (previously 20, then 4): under + // the paced scheme rounds are shallow + // and outer_tol exits early, so well-preconditioned methods use a handful of + // rounds while weakly-preconditioned ones get room to keep descending + // (native_ill CholQR2 genuinely uses 50) instead of being budget-truncated. + g_ir_n_steps = (int)opt_long(14, 50); + if (g_ir_n_steps < 1) { + std::cerr << "Error: ir_n_steps must be >= 1.\n"; + return 1; + } + // Outer early exit (structure unification with the Toeplitz pcg_ne): + // refinement stops once ||b - Jx||/||b|| meets this, capped at + // ir_n_steps. < 0 keeps the default 10*eps of the solve precision (the + // "refine until done" reading); 0 disables the check (always run all steps). + g_ir_outer_tol = opt_double(15, -1.0); + // Positions beyond 15: reject rather than ignore, so a stale job script fails + // loudly instead of silently running a different experiment than it encodes. + if (argc > dfactor_idx + 16) { + std::cerr << "Error: [ir_warm_start]/[bp_warm_start] CLI knobs were removed " + "(Blendenpik-only warm start, both variants always run). " + "Regenerate the job scripts.\n"; + return 1; + } + + if (mode == "irlsq_reg" && sparse_mode) { + std::cerr << "Error: mode 'irlsq_reg' is FEM-only; sparse input is not supported.\n"; + return 1; + } + + if (mode == "rspec") { + if (sparse_mode) { + std::cerr << "Error: mode 'rspec' is FEM-only; sparse input is not supported.\n"; + return 1; + } + if (power_j < 1 || power_j > 3) { + std::cerr << "Error: power_j must be in {1, 2, 3}; got " << power_j << "\n"; + return 1; + } + } + + std::cout << "=== CQRRT linop benchmark ===\n"; + std::cout << " mode: " << mode << "\n"; + if (sparse_mode) { + std::cout << " Input mode: sparse (single-matrix SparseLinOp)\n" + << " A file: " << A_file << "\n"; + } else { + std::cout << " Input mode: FEM composite (J = L^{-1} K V with L = chol(M))\n" + << " K file: " << K_file << "\n" + << " M file: " << M_file << "\n" + << " V file: " << V_file << "\n"; + } + std::cout << " d_factor: " << d_factor << "\n" + << " sketch_nnz: " << sketch_nnz << "\n" + << " block_size: " << block_size << "\n" + << " compute_cond: " << (compute_cond ? "yes" : "no") << "\n" + << " method_mask: " << method_mask << " ("; + // Full decoded roster: an echo limited to bits 0-4 would leave a mask-127 + // job log unable to show that Blendenpik/refine rows were selected. + // rspec has no Blendenpik dispatch (see run_rspec_benchmark's own decode + // call); echoing with_blendenpik=true there would print a roster the run + // then silently skips. + { + const bool echo_with_bp = (mode != "rspec"); + auto echo_algs = decode_method_mask(method_mask, echo_with_bp); + for (size_t i = 0; i < echo_algs.size(); ++i) + std::cout << (i ? " " : "") << echo_algs[i]; + } + std::cout << ")\n" + << " noise_level: " << noise_level << "\n" + << " omega: " << omega << "\n" + << " power_j: " << power_j << "\n" + << " num_runs: " << num_runs << "\n" +#ifdef _OPENMP + << " OpenMP threads: " << omp_get_max_threads() << "\n\n"; +#else + << " OpenMP threads: 1\n\n"; +#endif + + // ================================================================ + // Sparse mode: SparseLinOp directly, no Cholesky. + // ================================================================ + if (sparse_mode) { + int64_t m, n, nnz_A; + auto A_csr = load_csr_verbose("A", A_file, m, n, nnz_A); + RandLAPACK::linops::SparseLinOp> A_linop(m, n, A_csr); + + if (m < n) { + std::cerr << "Error: matrix must be overdetermined (m >= n), got " << m << "x" << n << "\n"; + return 1; + } + + // Sparse irlsq b construction: x_true ~ U(-1,1)^n, b = A x_true + scaled Gaussian noise. + std::vector x_true(n, (T)0); + { + std::mt19937 rng(42); + std::uniform_real_distribution dist((T)-1.0, (T)1.0); + for (auto& v : x_true) v = dist(rng); + } + std::vector b_clean(m, (T)0), noise_vec(m, (T)0); + A_linop(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, x_true.data(), n, (T)0.0, b_clean.data(), m); + T b_clean_norm = blas::nrm2(m, b_clean.data(), 1); + std::mt19937 noise_rng(13); + std::normal_distribution N01(0, 1); + for (auto& v : noise_vec) v = N01(noise_rng); + T raw_noise_norm = blas::nrm2(m, noise_vec.data(), 1); + T scale = noise_level * b_clean_norm / raw_noise_norm; + std::vector b(m, (T)0); + for (int64_t i = 0; i < m; ++i) b[i] = b_clean[i] + scale * noise_vec[i]; + std::cout << "Synthetic LS problem: ||x_true|| = " << blas::nrm2(n, x_true.data(), 1) + << ", ||b|| = " << blas::nrm2(m, b.data(), 1) << "\n\n"; + + return run_benchmark_inner( + A_linop, m, n, nnz_A, output_dir, num_runs, + d_factor, sketch_nnz, block_size, + compute_cond, + method_mask, noise_level, + 0L /*chol_time_us*/, + "A (" + A_file + ")", A_file, + &b, &x_true); + } + + // ================================================================ + // irlsq_reg mode (FEM-only): regularized augmented-operator preconditioner + // with independent preconditioner / solve precisions. Loads + builds its own + // (kappa-scaled, cast-down) chains, so it intercepts before the plain FEM load. + // (argv[1]) is the SOLVE precision; precond precision is a CLI knob. + // ================================================================ + if (mode == "irlsq_reg") { + // The noise_level CLI slot silently defaults to 0.05. irlsq_reg's + // consistent-RHS backward-error reading is cleanest at noise_level=0, + // so a nonzero DEFAULTED value + // (as opposed to one the caller explicitly asked for) is worth flagging loudly + // rather than baking a silent, easy-to-miss assumption into the CSV. + if (!noise_level_explicit && noise_level != (T)0) { + std::cerr << "WARNING: irlsq_reg noise_level defaulted to " << (double)noise_level + << " (not explicitly set on the CLI). Pass noise_level=0 explicitly for " + "the pure consistent-RHS backward-error reading, or pass the intended " + "nonzero value explicitly to silence this warning.\n"; + } + std::string solve_prec_str = (sizeof(T) == 8) ? "double" : "single"; + std::cout << "\n=== IR-LSQ-reg mode (regularized augmented operator) ===\n" + << " kappa_target: " << kappa_target << "\n" + << " mu_factor: " << mu_factor << "\n" + << " noise_level: " << (double)noise_level << "\n" + << " precond_prec: " << precond_prec << "\n" + << " solve_prec: " << solve_prec_str << "\n\n"; + return dispatch_irlsq_reg(precond_prec, K_file, M_file, V_file, + output_dir, num_runs, (double)d_factor, sketch_nnz, block_size, + compute_cond, method_mask, kappa_target, mu_factor, (double)noise_level, solve_prec_str); + } + + // ================================================================ + // FEM mode: load K, V; Cholesky-factorize M; build J = L^{-1} K V. + // ================================================================ + int64_t m_K, n_K, nnz_K; + auto K_csr = load_csr_verbose("K (stiffness)", K_file, m_K, n_K, nnz_K); + if (m_K != n_K) { + std::cerr << "Error: K must be square; got " << m_K << " x " << n_K << "\n"; + return 1; + } + RandLAPACK::linops::SparseLinOp> K_op(m_K, m_K, K_csr); + + int64_t m_V, n_V, nnz_V; + auto V_csr = load_csr_verbose("V (prolongation)", V_file, m_V, n_V, nnz_V); + if (m_V != m_K) { + std::cerr << "Error: V row count (" << m_V << ") must match K size (" << m_K << ")\n"; + return 1; + } + if (m_V < n_V) { + std::cerr << "Error: need tall V (m_fine >= n_coarse); got " << m_V << " x " << n_V << "\n"; + return 1; + } + RandLAPACK::linops::SparseLinOp> V_op(m_V, n_V, V_csr); + + std::cout << "Factorizing M = L L^T from " << M_file << "... " << std::flush; + RandLAPACK_extras::linops::CholSolverLinOp L_inv_op(M_file, /*half_solve=*/true); + auto chol_start = steady_clock::now(); + L_inv_op.factorize(); + auto chol_stop = steady_clock::now(); + long chol_time_us = duration_cast(chol_stop - chol_start).count(); + std::cout << "done (" << chol_time_us << " us)\n"; + if (L_inv_op.n_rows != m_K) { + std::cerr << "Error: M size (" << L_inv_op.n_rows << ") must match K size (" + << m_K << ")\n"; + return 1; + } + + int64_t m = m_V; + int64_t n = n_V; + + // -------- RSPEC mode (Algorithm 4): build C = L^T X^{-1} L and V_app = C^j V_FEM. -------- + if (mode == "rspec") { + std::cout << "\n=== RSPEC mode (Algorithm 4) ===\n" + << " omega: " << omega << "\n" + << " power_j: " << power_j << "\n"; + + // Load M as a CSR so we can form X = K - omega*M via shared-pattern axpby. + int64_t m_M, n_M, nnz_M; + auto M_csr = load_csr_verbose("M (mass, for X assembly)", M_file, m_M, n_M, nnz_M); + if (m_M != m_K || n_M != m_K) { + std::cerr << "Error: M size (" << m_M << " x " << n_M + << ") must match K size " << m_K << "\n"; + return 1; + } + + // 1. X = K - omega * M (CSR, shares sparsity with K and M). + std::cout << "Forming X = K - omega*M ... " << std::flush; + auto X_csr = RandLAPACK_extras::sparse_axpby_shared_pattern( + (T)1.0, K_csr, -(T)omega, M_csr); + std::cout << "done (nnz=" << X_csr.nnz << ")\n"; + + // 2. Factor X via sparse Cholesky. + // + // X = K - omega*M is SPD as long as omega < lambda_min(K, M) (the smallest + // generalized eigenvalue of the (K, M) pencil). For omega = 0 and the + // near-zero shifts this application uses, X stays positive definite, so a + // sparse Cholesky factorization suffices: it confines Eigen to the + // factorization and applies X^{-1} via RandBLAS sparse TRSM (CholSolverLinOp). + // An interior shift (omega >= lambda_min) would make X indefinite; Cholesky + // would then (correctly) fail and an indefinite solver would be needed. + std::cout << "Factorizing X = L L^T (sparse Cholesky) ... " << std::flush; + RandLAPACK_extras::linops::CholSolverLinOp X_inv_op(X_csr, /*half_solve=*/false); + auto x_fact_start = steady_clock::now(); + try { + X_inv_op.factorize(); + } catch (RandBLAS::Error const& e) { + auto x_fact_stop = steady_clock::now(); + long x_fact_us = duration_cast(x_fact_stop - x_fact_start).count(); + std::cerr << "\nCholesky factorization of X failed (X not SPD, omega at/above " + "lambda_min, or near an eigenvalue): " << e.what() << "\n"; + + // Write a single sentinel row to the CSV and return cleanly. + std::string results_file = output_dir + "/" + make_run_timestamp() + "_rspec_results.csv"; + + std::vector> stub; + rspec_result r{}; + r.m = m_K; r.n = n_V; + r.run_idx = 0; + r.alg_name = "factorize_failed"; + r.qr_status = -99; + r.qr_time_us = -1; + r.peak_rss_kb = -1; // -1 sentinel: never measured, not 0 + r.analytical_kb = -1; + r.factor_time_us = x_fact_us; + r.rspec_total_us = x_fact_us; + r.orth_error = std::numeric_limits::quiet_NaN(); + int top_k = (int)std::min(10, n_V); + r.top_eigvals.assign(top_k, std::numeric_limits::quiet_NaN()); + r.top_residuals.assign(top_k, std::numeric_limits::quiet_NaN()); + stub.push_back(r); + write_rspec_csv(results_file, stub, m_K, n_V, num_runs, + K_file, M_file, V_file, omega, power_j, + sketch_nnz, block_size, method_mask, top_k); + std::cout << "Stub CSV written to " << results_file << " (qr_status=-99).\n"; + return 0; + } + auto x_fact_stop = steady_clock::now(); + long x_factor_time_us = duration_cast(x_fact_stop - x_fact_start).count(); + std::cout << "done (" << x_factor_time_us << " us)\n"; + + // 4. L_op: wrap the L = chol(M) factor as a sparse linop (non-owning view). + auto L_csc = L_inv_op.make_L_csc(); + RandLAPACK::linops::SparseLinOp> L_op(m_K, m_K, L_csc); + + // 5. Compose C = L^T * X^{-1} * L. + RandLAPACK::linops::TransposedOp L_T_op(L_op); + RandLAPACK::linops::CompositeOperator inner_op(m_K, m_K, X_inv_op, L_op); + inner_op.block_size = block_size; + RandLAPACK::linops::CompositeOperator C_op(m_K, m_K, L_T_op, inner_op); + C_op.block_size = block_size; + + // 6. V_app = C^j * V_FEM (implicit). + RandLAPACK::linops::PowerOp Cj_op(C_op, (int)power_j); + RandLAPACK::linops::CompositeOperator V_app_op(m_K, n_V, Cj_op, V_op); + V_app_op.block_size = block_size; + + std::cout << "Operator chain: V_app = C^" << power_j + << " * V_FEM (" << m_K << " x " << n_V << ")\n\n"; + + long total_factor_us = chol_time_us + x_factor_time_us; + return run_rspec_benchmark( + V_app_op, C_op, + m_K, n_V, output_dir, num_runs, + d_factor, sketch_nnz, block_size, + method_mask, total_factor_us, + K_file, M_file, V_file, omega, power_j); + } + + RandLAPACK::linops::CompositeOperator KV_op(m, n, K_op, V_op); + KV_op.block_size = block_size; + RandLAPACK::linops::CompositeOperator J_op(m, n, L_inv_op, KV_op); + J_op.block_size = block_size; + std::cout << "Composite operator J = L^{-1} K V : " << m << " x " << n << "\n\n"; + + // FEM irlsq b construction: b = L^{-1} * r, r ~ N(0, 1)^{m_K}. No ground truth x_true. + std::vector r(m_K, (T)0); + { + std::mt19937 rng_b(13); + std::normal_distribution N01(0, 1); + for (auto& v : r) v = N01(rng_b); + } + std::vector b(m_K, (T)0); + L_inv_op(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m_K, 1, m_K, (T)1.0, r.data(), m_K, (T)0.0, b.data(), m_K); + std::cout << "FEM IR-LSQ b: ||b|| = " << blas::nrm2(m_K, b.data(), 1) + << " (b = L^{-1} r, r ~ N(0,1)^M)\n\n"; + + return run_benchmark_inner( + J_op, m, n, nnz_K, output_dir, num_runs, + d_factor, sketch_nnz, block_size, + compute_cond, + method_mask, noise_level, + chol_time_us, + "L^{-1} K V (M=" + M_file + ")", K_file, + &b, nullptr /*x_true_ptr: FEM has no ground truth*/); +} + +int main(int argc, char* argv[]) { + if (argc < 2) { + std::cerr << "Usage: " << argv[0] + << " \n" + << " sparse mode: 'sparse' " + << " [sketch_nnz] [block_size] [compute_cond] [method_mask] [noise_level]\n" + << " FEM mode: " + << " [sketch_nnz] [block_size] [compute_cond] [method_mask] [noise_level]" + << " [omega] [power_j]\n" + << " mode = irlsq | rspec | irlsq_reg (rspec/irlsq_reg are FEM-only)\n"; + return 1; + } + + std::string precision = argv[1]; + if (precision == "double") { + return run_benchmark(argc, argv); + } else if (precision == "float" || precision == "single") { + return run_benchmark(argc, argv); + } else { + std::cerr << "Unknown precision: " << precision << " (use 'double'/'float'/'single')\n"; + return 1; + } +} diff --git a/benchmark/bench_CQRRT_linops/CQRRT_linop_basic.cc b/benchmark/bench_CQRRT_linops/CQRRT_linop_basic.cc index 7b8eb0145..9fbfe84d7 100644 --- a/benchmark/bench_CQRRT_linops/CQRRT_linop_basic.cc +++ b/benchmark/bench_CQRRT_linops/CQRRT_linop_basic.cc @@ -15,6 +15,7 @@ #include #include #include +#include #ifdef _OPENMP #include #endif @@ -22,9 +23,9 @@ // Extras utilities for Matrix Market I/O #include "../../extras/misc/ext_util.hh" #include "RandLAPACK/testing/rl_test_utils.hh" +#include "cqrrt_bench_common.hh" // Linops algorithms (now in main RandLAPACK) -#include "rl_cqrrt_linops.hh" #include "rl_cholqr_linops.hh" #include "rl_scholqr3_linops.hh" #include "RandLAPACK/testing/rl_memory_tracker.hh" @@ -33,16 +34,40 @@ using std::chrono::steady_clock; using std::chrono::duration_cast; using std::chrono::microseconds; -// Common quality + timing fields shared by all algorithms +// Shared helpers (cqrrt_bench_common.hh): env-knob provenance, the +// per-pass Cholesky-shift fold, and argv/host CSV provenance. +using RandLAPACK::bench::write_env_line; +using RandLAPACK::bench::fold_chol_shift; +using RandLAPACK::bench::quote_join_argv; +using RandLAPACK::bench::get_hostname; + +// Common quality + timing fields shared by all algorithms. Every field +// defaults to a -1/failure sentinel (never a silently-perfect 0) so a +// skipped or failed algorithm reads unambiguously in the CSV. template struct alg_quality { - T orth_error; // ||Q^T Q - I|| / sqrt(n) - bool is_orthonormal; // Is full Q block orthonormal? - int64_t max_orth_cols; // Maximum orthonormal prefix - long time; // Total computation time (microseconds) - long peak_rss_kb; // Peak RSS increase during algorithm call (KB) - long analytical_kb; // Analytical peak working memory (KB) - std::vector breakdown; // Per-subroutine timings (excludes total) + int qr_status = -1; // 0 = success; driver's own code on failure; + // -1 = not yet run (skip_dense, or a bug if + // still -1 after run_algorithms returns) + T orth_error = (T)-1; // ||Q^T Q - I|| / sqrt(n) + bool is_orthonormal = false; // Is full Q block orthonormal? + int64_t max_orth_cols = -1; // Maximum orthonormal prefix + long time = -1; // Total computation time (microseconds) + long peak_rss_kb = -1; // Peak RSS increase during algorithm call (KB) + long analytical_kb = -1; // Analytical peak working memory (KB) + std::vector breakdown; // Per-subroutine timings (excludes total); + // always sized to the algorithm's fixed + // slot count (10/5/17/10) regardless of + // success, failure, or skip_dense, so the + // breakdown CSV row is always full-width. + int chol_retries = -1; // -1 = the row failed, or was skipped + // (dense_cqrrt under skip_dense=1); every + // algorithm in this file, including dense + // CQRRT_expl, has an adaptive Cholesky-shift + // retry mechanism, so -1 no longer means + // "no mechanism" + T chol_shift_abs = (T)-1; + T chol_shift_rel = (T)-1; }; template @@ -68,15 +93,13 @@ template static void compute_Q_from_R( GLO& A_op, T* R, int64_t ldr, T* Q_out, int64_t m, int64_t n) { - // Step 1: Materialize A into Q_out: Q_out = A * I T* Eye = new T[n * n](); RandLAPACK::util::eye(n, n, Eye); A_op(Side::Left, Layout::ColMajor, Op::NoTrans, Op::NoTrans, m, n, n, (T)1.0, Eye, n, (T)0.0, Q_out, m); - delete[] Eye; - // Step 2: Solve Q * R = A for Q via trsm (backward stable, no explicit inverse) blas::trsm(Layout::ColMajor, Side::Right, Uplo::Upper, Op::NoTrans, Diag::NonUnit, m, n, (T)1.0, R, ldr, Q_out, m); + delete[] Eye; } // Core algorithm runner: operates on a pre-constructed SparseLinOp. @@ -108,110 +131,134 @@ static std::vector> run_algorithms( T tol = std::pow(std::numeric_limits::epsilon(), 0.85); // Single reusable Q buffer for uniform Q = A * R^{-1} computation across all algorithms - std::vector Q_uniform(m * n); + T* Q_uniform = new T[m * n]; + + // RSS window (unified): every algorithm's peak_rss_kb is measured around + // its OWN single timed call() below, all with test_mode=false. CQRRT_linop + // and CholQR used to run a SEPARATE untimed pre-call for the RSS + // measurement, whose stated rationale ("excludes the test_mode=true + // Q-factor allocation") never applied, since the timed call was already + // test_mode=false, so that extra pass measured nothing the timed call's + // own window wouldn't. Dropped in favor of the single-window scheme + // sCholQR3 already used. // ============================================================ // Run CQRRT (preconditioned Cholesky QR) - multiple runs // ============================================================ - // Peak RSS measured separately with test_mode=false to exclude Q-factor allocation. - // With column-blocking, test_mode reallocates A_pre from m*b_eff to m*n for Q. { - // RSS measurement (test_mode=false) - long cqrrt_peak_rss_kb = 0; - { - std::vector R_rss(n * n, 0.0); - auto state_rss = run_states[0]; - RandLAPACK::CQRRT_linops CQRRT_rss(false, tol, false); - CQRRT_rss.nnz = sketch_nnz; - CQRRT_rss.block_size = block_size; - RandLAPACK::PeakRSSTracker cqrrt_mem; - cqrrt_mem.start(); - CQRRT_rss.call(A_linop, R_rss.data(), n, d_factor, state_rss); - cqrrt_peak_rss_kb = cqrrt_mem.stop(); - } - + T* R_cqrrt = new T[n * n]; for (int64_t run = 0; run < num_runs; ++run) { - std::vector R_cqrrt(n * n, 0.0); + std::fill(R_cqrrt, R_cqrrt + n * n, (T)0); auto state_copy = run_states[run]; // Per-run RNG state RandLAPACK::CQRRT_linops CQRRT_QR(true, tol, false); // timing=true, test_mode=false CQRRT_QR.nnz = sketch_nnz; CQRRT_QR.block_size = block_size; - CQRRT_QR.call(A_linop, R_cqrrt.data(), n, d_factor, state_copy); - results[run].cqrrt.time = CQRRT_QR.times[10]; // total_t_dur - results[run].cqrrt.peak_rss_kb = cqrrt_peak_rss_kb; - results[run].cqrrt.breakdown.assign(CQRRT_QR.times.begin(), CQRRT_QR.times.begin() + 10); - - // Uniform Q computation for every run: Q = A * R^{-1} via operator - compute_Q_from_R(A_linop, R_cqrrt.data(), n, Q_uniform.data(), m, n); - results[run].cqrrt.orth_error = RandLAPACK::testing::orthogonality_error(Q_uniform.data(), m, n); - results[run].cqrrt.is_orthonormal = (results[run].cqrrt.orth_error <= std::pow(std::numeric_limits::epsilon(), (T)0.75)); - results[run].cqrrt.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform.data(), m, n); + RandLAPACK::PeakRSSTracker cqrrt_mem; + cqrrt_mem.start(); + int status = CQRRT_QR.call(A_linop, R_cqrrt, n, d_factor, state_copy); + results[run].cqrrt.peak_rss_kb = cqrrt_mem.stop(); + results[run].cqrrt.qr_status = status; + + if (status == 0) { + results[run].cqrrt.time = CQRRT_QR.total_us(); + results[run].cqrrt.breakdown.assign(CQRRT_QR.times.begin(), CQRRT_QR.times.end() - 1); + results[run].cqrrt.chol_retries = CQRRT_QR.n_chol_retries; + fold_chol_shift(results[run].cqrrt.chol_shift_abs, results[run].cqrrt.chol_shift_rel, + CQRRT_QR.chol_applied_shifts, CQRRT_QR.chol_gram_traces, 1); + + // Uniform Q computation for every run: Q = A * R^{-1} via operator + compute_Q_from_R(A_linop, R_cqrrt, n, Q_uniform, m, n); + results[run].cqrrt.orth_error = RandLAPACK::testing::orthogonality_error(Q_uniform, m, n); + results[run].cqrrt.is_orthonormal = (results[run].cqrrt.orth_error <= std::pow(std::numeric_limits::epsilon(), (T)0.75)); + results[run].cqrrt.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform, m, n); + } else { + // Failure: no valid R, so no Q and no orthogonality metrics. + // Full-width -1 breakdown (never a truncated/empty vector). + results[run].cqrrt.breakdown.assign(10, -1L); + std::cerr << "Warning: CQRRT_linop call() failed (run=" << run + << ", status=" << status << "); quality/timing fields set to -1.\n"; + } } + delete[] R_cqrrt; } // ============================================================ // Run CholQR (unpreconditioned Cholesky QR) - multiple runs // ============================================================ - // Peak RSS measured separately with test_mode=false to exclude Q-factor allocation. - // With column-blocking, test_mode reallocates A_temp from m*b_eff to m*n for Q. { - // RSS measurement (test_mode=false) - long cholqr_peak_rss_kb = 0; - { - std::vector R_rss(n * n, 0.0); - RandLAPACK::CholQR_linops CholQR_rss(false, tol, false); - CholQR_rss.block_size = block_size; - RandLAPACK::PeakRSSTracker cholqr_mem; - cholqr_mem.start(); - CholQR_rss.call(A_linop, R_rss.data(), n); - cholqr_peak_rss_kb = cholqr_mem.stop(); - } - + T* R_cholqr = new T[n * n]; for (int64_t run = 0; run < num_runs; ++run) { - std::vector R_cholqr(n * n, 0.0); + std::fill(R_cholqr, R_cholqr + n * n, (T)0); RandLAPACK::CholQR_linops CholQR_alg(true, tol, false); // timing=true, test_mode=false CholQR_alg.block_size = block_size; - CholQR_alg.call(A_linop, R_cholqr.data(), n); - results[run].cholqr.time = CholQR_alg.times[5]; // total - results[run].cholqr.peak_rss_kb = cholqr_peak_rss_kb; - results[run].cholqr.breakdown.assign(CholQR_alg.times.begin(), CholQR_alg.times.begin() + 5); - - // Uniform Q computation for every run - compute_Q_from_R(A_linop, R_cholqr.data(), n, Q_uniform.data(), m, n); - results[run].cholqr.orth_error = RandLAPACK::testing::orthogonality_error(Q_uniform.data(), m, n); - results[run].cholqr.is_orthonormal = (results[run].cholqr.orth_error <= std::pow(std::numeric_limits::epsilon(), (T)0.75)); - results[run].cholqr.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform.data(), m, n); + RandLAPACK::PeakRSSTracker cholqr_mem; + cholqr_mem.start(); + int status = CholQR_alg.call(A_linop, R_cholqr, n); + results[run].cholqr.peak_rss_kb = cholqr_mem.stop(); + results[run].cholqr.qr_status = status; + + if (status == 0) { + results[run].cholqr.time = CholQR_alg.total_us(); + results[run].cholqr.breakdown.assign(CholQR_alg.times.begin(), CholQR_alg.times.end() - 1); + results[run].cholqr.chol_retries = CholQR_alg.n_chol_retries; + fold_chol_shift(results[run].cholqr.chol_shift_abs, results[run].cholqr.chol_shift_rel, + CholQR_alg.chol_applied_shifts, CholQR_alg.chol_gram_traces, 1); + + // Uniform Q computation for every run + compute_Q_from_R(A_linop, R_cholqr, n, Q_uniform, m, n); + results[run].cholqr.orth_error = RandLAPACK::testing::orthogonality_error(Q_uniform, m, n); + results[run].cholqr.is_orthonormal = (results[run].cholqr.orth_error <= std::pow(std::numeric_limits::epsilon(), (T)0.75)); + results[run].cholqr.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform, m, n); + } else { + results[run].cholqr.breakdown.assign(5, -1L); + std::cerr << "Warning: CholQR call() failed (run=" << run + << ", status=" << status << "); quality/timing fields set to -1.\n"; + } } + delete[] R_cholqr; } // ============================================================ // Run sCholQR3 (shifted Cholesky QR with 3 iterations) - multiple runs // ============================================================ { + T* R_scholqr3 = new T[n * n]; for (int64_t run = 0; run < num_runs; ++run) { - std::vector R_scholqr3(n * n, 0.0); + std::fill(R_scholqr3, R_scholqr3 + n * n, (T)0); RandLAPACK::sCholQR3_linops sCholQR3_alg(true, tol, false); // timing=true, test_mode=false sCholQR3_alg.block_size = block_size; RandLAPACK::PeakRSSTracker scholqr3_mem; scholqr3_mem.start(); - sCholQR3_alg.call(A_linop, R_scholqr3.data(), n); + int status = sCholQR3_alg.call(A_linop, R_scholqr3, n); results[run].scholqr3.peak_rss_kb = scholqr3_mem.stop(); - - results[run].scholqr3.time = sCholQR3_alg.times[12]; // total - results[run].scholqr3.breakdown.assign(sCholQR3_alg.times.begin(), sCholQR3_alg.times.begin() + 12); - - // Uniform Q computation (same as all other algorithms) - compute_Q_from_R(A_linop, R_scholqr3.data(), n, Q_uniform.data(), m, n); - results[run].scholqr3.orth_error = RandLAPACK::testing::orthogonality_error(Q_uniform.data(), m, n); - results[run].scholqr3.is_orthonormal = (results[run].scholqr3.orth_error <= std::pow(std::numeric_limits::epsilon(), (T)0.75)); - results[run].scholqr3.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform.data(), m, n); + results[run].scholqr3.qr_status = status; + + if (status == 0) { + results[run].scholqr3.time = sCholQR3_alg.total_us(); + // breakdown (17): alloc, fwd1, adj1, chol1, upd1, fwd2, adj2, gemm2, chol2, upd2, fwd3, adj3, gemm3, chol3, upd3, q_mat, rest + results[run].scholqr3.breakdown.assign(sCholQR3_alg.times.begin(), sCholQR3_alg.times.end() - 1); + results[run].scholqr3.chol_retries = sCholQR3_alg.n_chol_retries; + fold_chol_shift(results[run].scholqr3.chol_shift_abs, results[run].scholqr3.chol_shift_rel, + sCholQR3_alg.chol_applied_shifts, sCholQR3_alg.chol_gram_traces, 3); + + // Uniform Q computation (same as all other algorithms) + compute_Q_from_R(A_linop, R_scholqr3, n, Q_uniform, m, n); + results[run].scholqr3.orth_error = RandLAPACK::testing::orthogonality_error(Q_uniform, m, n); + results[run].scholqr3.is_orthonormal = (results[run].scholqr3.orth_error <= std::pow(std::numeric_limits::epsilon(), (T)0.75)); + results[run].scholqr3.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform, m, n); + } else { + results[run].scholqr3.breakdown.assign(17, -1L); + std::cerr << "Warning: sCholQR3 call() failed (run=" << run + << ", status=" << status << "); quality/timing fields set to -1.\n"; + } } + delete[] R_scholqr3; } // ============================================================ @@ -220,13 +267,14 @@ static std::vector> run_algorithms( // Peak RSS with compute_Q=true is correct: Q overwrites A_materialized in-place (no extra allocation). // Skipped when skip_dense=true (e.g., for large file-input matrices where m*n dense doesn't fit in memory). if (!skip_dense) { + T* I_mat = new T[n * n](); + RandLAPACK::util::eye(n, n, I_mat); + T* R_dense = new T[n * n]; for (int64_t run = 0; run < num_runs; ++run) { RandLAPACK::PeakRSSTracker dense_mem; dense_mem.start(); // Step 1: Materialize the operator by multiplying with identity - T* I_mat = new T[n * n](); - RandLAPACK::util::eye(n, n, I_mat); T* A_materialized = new T[m * n](); auto materialize_start = steady_clock::now(); @@ -235,51 +283,80 @@ static std::vector> run_algorithms( auto materialize_stop = steady_clock::now(); long materialize_time = duration_cast(materialize_stop - materialize_start).count(); - delete[] I_mat; - // Step 2: Call rl_cqrrt with timing, Q-factor disabled (computed uniformly below) // Uses same per-run RNG state as CQRRT_linop for fair comparison - std::vector R_dense(n * n, 0.0); + std::fill(R_dense, R_dense + n * n, (T)0); auto state_copy = run_states[run]; // Same RNG state as CQRRT_linop's run RandLAPACK::CQRRT dense_alg(true, tol); // timing=true dense_alg.compute_Q = false; dense_alg.orthogonalization = false; dense_alg.nnz = sketch_nnz; - dense_alg.call(m, n, A_materialized, m, R_dense.data(), n, d_factor, state_copy); + int status = dense_alg.call(m, n, A_materialized, m, R_dense, n, d_factor, state_copy); results[run].dense_cqrrt.peak_rss_kb = dense_mem.stop(); + results[run].dense_cqrrt.qr_status = status; delete[] A_materialized; // No longer needed (Q computed via operator) - // Total = materialization + algorithm total (Q excluded from algo total) - results[run].dense_cqrrt.time = materialize_time + dense_alg.times[9]; - // Breakdown matches linop CQRRT layout: materialize, saso, qr, trtri(=0), precond, gram, trmm_gram(=0), potrf, finalize, rest - results[run].dense_cqrrt.breakdown = { - materialize_time, - dense_alg.times[0], // saso - dense_alg.times[1], // qr - 0L, // trtri (always 0 for dense) - dense_alg.times[3], // precond - dense_alg.times[4], // gram - 0L, // trmm_gram (always 0 for dense) - dense_alg.times[6], // potrf - dense_alg.times[7], // finalize - dense_alg.times[8], // rest - }; - - // Uniform Q computation for every run - compute_Q_from_R(A_linop, R_dense.data(), n, Q_uniform.data(), m, n); - results[run].dense_cqrrt.orth_error = RandLAPACK::testing::orthogonality_error(Q_uniform.data(), m, n); - results[run].dense_cqrrt.is_orthonormal = (results[run].dense_cqrrt.orth_error <= std::pow(std::numeric_limits::epsilon(), (T)0.75)); - results[run].dense_cqrrt.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform.data(), m, n); + if (status == 0) { + // Total = materialization + algorithm total (Q excluded from algo total) + results[run].dense_cqrrt.time = materialize_time + dense_alg.total_us(); + // Breakdown matches linop CQRRT layout: materialize, saso, qr, trtri(=0), precond, gram, trmm_gram(=0), potrf, finalize, rest + results[run].dense_cqrrt.breakdown = { + materialize_time, + dense_alg.times[0], // saso + dense_alg.times[1], // qr + 0L, // trtri (always 0 for dense) + dense_alg.times[3], // precond + dense_alg.times[4], // gram + 0L, // trmm_gram (always 0 for dense) + dense_alg.times[6], // potrf + dense_alg.times[7], // finalize + dense_alg.times[8], // rest + }; + // Dense CQRRT (rl_cqrrt.hh CQRRT, not CQRRT_linops) has + // the same adaptive Cholesky-shift retry as the other three + // algorithms; fold its record the same + // way the other rows do (single shift-record entry, npasses=1, + // same as CQRRT_linops's own single-pass record). + results[run].dense_cqrrt.chol_retries = dense_alg.n_chol_retries; + fold_chol_shift(results[run].dense_cqrrt.chol_shift_abs, results[run].dense_cqrrt.chol_shift_rel, + dense_alg.chol_applied_shifts, dense_alg.chol_gram_traces, 1); + + // Uniform Q computation for every run + compute_Q_from_R(A_linop, R_dense, n, Q_uniform, m, n); + results[run].dense_cqrrt.orth_error = RandLAPACK::testing::orthogonality_error(Q_uniform, m, n); + results[run].dense_cqrrt.is_orthonormal = (results[run].dense_cqrrt.orth_error <= std::pow(std::numeric_limits::epsilon(), (T)0.75)); + results[run].dense_cqrrt.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform, m, n); + } else { + results[run].dense_cqrrt.breakdown.assign(10, -1L); + std::cerr << "Warning: CQRRT_expl (dense) call() failed (run=" << run + << ", status=" << status << "); quality/timing fields set to -1.\n"; + } + } + delete[] I_mat; + delete[] R_dense; + } else { + // skip_dense=true: CQRRT_expl never runs. Every field carries the -1 + // sentinel (never a value-initialized 0, which would misreport as + // "measured, perfectly orthogonal") so a reader can tell "skipped" + // apart from "ran and was flawless". qr_status=-1 marks + // "not run"; it is deliberately not one of CQRRT's own nonzero + // failure codes, which mean the algorithm genuinely tried and failed. + for (int64_t r = 0; r < num_runs; ++r) { + results[r].dense_cqrrt.breakdown.assign(10, -1L); + results[r].dense_cqrrt.qr_status = -1; } - } // if (!skip_dense) + } - // Compute analytical peak working memory for each algorithm (same for all runs) + // Compute analytical peak working memory for each algorithm (same for all runs). + // dense_cqrrt_akb is gated the same way as the rest of the dense row: + // skip_dense=true means CQRRT_expl never ran, so its analytical model is + // -1 ("not applicable"), not a number nobody measured against. long cqrrt_akb = RandLAPACK::cqrrt_linops_analytical_kb(m, n, d_factor, block_size); long cholqr_akb = RandLAPACK::cholqr_linops_analytical_kb(m, n, block_size); long scholqr3_akb = RandLAPACK::scholqr3_linops_analytical_kb(m, n, block_size); - long dense_cqrrt_akb = RandLAPACK::dense_cqrrt_analytical_kb(m, n, d_factor); + long dense_cqrrt_akb = skip_dense ? -1L : RandLAPACK::dense_cqrrt_analytical_kb(m, n, d_factor); for (int64_t r = 0; r < num_runs; ++r) { results[r].cqrrt.analytical_kb = cqrrt_akb; results[r].cholqr.analytical_kb = cholqr_akb; @@ -287,6 +364,7 @@ static std::vector> run_algorithms( results[r].dense_cqrrt.analytical_kb = dense_cqrrt_akb; } + delete[] Q_uniform; return results; } @@ -354,12 +432,9 @@ static std::vector> run_single_test_from_file( } // Load matrix from Matrix Market file - auto A_coo = RandLAPACK_extras::coo_from_matrix_market(filename); - int64_t m = A_coo.n_rows; - int64_t n = A_coo.n_cols; - T actual_density = static_cast(A_coo.nnz) / (static_cast(m) * n); - RandBLAS::sparse_data::csr::CSRMatrix A_csr(m, n); - RandBLAS::sparse_data::conversions::coo_to_csr(A_coo, A_csr); + int64_t m, n, nnz; + auto A_csr = load_csr(filename, m, n, nnz); + T actual_density = static_cast(nnz) / (static_cast(m) * n); RandLAPACK::linops::SparseLinOp> A_linop(m, n, A_csr); T cond_num = std::numeric_limits::quiet_NaN(); @@ -384,7 +459,7 @@ static void write_csv_headers( std::ofstream& out, std::ofstream& breakdown, const std::string& precision, double d_factor, int64_t sketch_nnz, int64_t block_size, int64_t num_runs, int num_threads, - const std::string& extra_comment); + const std::string& extra_comment, const std::string& argv_line); static void prepend_runtime(const std::string& filepath, double seconds); template @@ -406,7 +481,7 @@ static int run_benchmark(int argc, char *argv[]) { std::cerr << " density : Target density (e.g., 0.1); bandwidth derived as round(density*n - 1)" << std::endl; std::cerr << " d_factor : Sketching dimension factor for CQRRT_linop (e.g., 2.0)" << std::endl; std::cerr << " sketch_nnz : (Optional) Nonzeros per column in SASO sketch (default: 4)" << std::endl; - std::cerr << " block_size : (Optional) Column-block size for CQRRT_linop/CholQR/sCholQR3 Gram (0 = full, default: 0)" << std::endl; + std::cerr << " block_size : (Optional) Column-block size for CQRRT_linop/CholQR/sCholQR3 Gram (0 = full, default: 256 = paper b)" << std::endl; std::cerr << "\nExample:" << std::endl; std::cerr << " " << argv[0] << " double ./output 30 3 1000 30000 100 1e9 0.05 2.0 4 100" << std::endl; std::cerr << " (Tests 30 matrices from 1000x10 to 30000x300, aspect ratio 100:1, κ=1e9, density≈0.05, 3 runs each)" << std::endl; @@ -416,10 +491,10 @@ static int run_benchmark(int argc, char *argv[]) { // Parse arguments std::string precision = argv[1]; std::string output_dir = argv[2]; - int64_t num_sizes = std::stol(argv[3]); - int64_t num_runs = std::stol(argv[4]); - int64_t m_start = std::stol(argv[5]); - int64_t m_end = std::stol(argv[6]); + int64_t num_sizes = std::stoll(argv[3]); + int64_t num_runs = std::stoll(argv[4]); + int64_t m_start = std::stoll(argv[5]); + int64_t m_end = std::stoll(argv[6]); double aspect_ratio = std::stod(argv[7]); double cond_num = std::stod(argv[8]); double density = std::stod(argv[9]); @@ -427,9 +502,24 @@ static int run_benchmark(int argc, char *argv[]) { // Default sketch_nnz=4: the Givens-based matrix generator produces // high-coherence matrices (non-uniform leverage scores), so nnz >= 4 // is needed for reliable SASO sketching (nnz=2 causes sporadic spikes). - int64_t sketch_nnz = (argc >= 12) ? std::stol(argv[11]) : 4; - int64_t block_size = (argc >= 13) ? std::stol(argv[12]) : 0; + int64_t sketch_nnz = (argc >= 12) ? std::stoll(argv[11]) : 4; + // Default block_size=256 matches the paper's b=256 (blocked Gram, tall + // intermediate never fully materialized); pass 0 explicitly for unblocked. + int64_t block_size = (argc >= 13) ? std::stoll(argv[12]) : 256; + + // Loud validation: num_sizes=0 leaves `sizes` empty and + // sizes.front()/.back() below is UB; num_runs<=0 leaves `results` empty + // and every per-run write is UB (num_runs=0 was silently accepted before). + if (num_sizes < 1) { + std::cerr << "Error: num_sizes must be >= 1 (got " << num_sizes << ")\n"; + return 1; + } + if (num_runs < 1) { + std::cerr << "Error: num_runs must be >= 1 (got " << num_runs << ")\n"; + return 1; + } + std::string argv_line = quote_join_argv(argc, argv); auto benchmark_start = steady_clock::now(); // Generate date/time prefix @@ -482,7 +572,7 @@ static int run_benchmark(int argc, char *argv[]) { << "# Condition number: " << cond_num << "\n" << "# Target density: " << density << "\n"; write_csv_headers(out, breakdown, precision, d_factor, - sketch_nnz, block_size, num_runs, num_threads, extra.str()); + sketch_nnz, block_size, num_runs, num_threads, extra.str(), argv_line); // Warmup run to trigger library initialization (MKL thread pools, memory allocators, etc.) // This ensures first reported iteration has accurate memory measurements. @@ -539,6 +629,35 @@ static void write_results_to_csv( const alg_quality* algos[] = {&result.cqrrt, &result.cholqr, &result.scholqr3, &result.dense_cqrrt}; + // Full-width guard: every algorithm's breakdown must be exactly its + // fixed slot count (10/5/17/10) on every path (success, driver + // failure, or skip_dense), or the breakdown row drifts out of + // alignment with the header. run_algorithms guarantees this by + // construction; a short vector here is not a corrupted campaign, so + // it is padded (with the same -1 sentinel the failure paths use) and + // warned about rather than discarding a finished run. A vector LONGER + // than the header width cannot be safely truncated without silently + // dropping real data, so that case still throws: it means something + // upstream is producing more fields than the schema declares. + auto fit_breakdown = [](std::vector v, size_t expected, const char* name) { + if (v.size() > expected) { + randlapack_require(false) + << name << " breakdown width mismatch: expected " << expected + << " fields, got " << v.size() << " (exceeds the header width, " + << "which indicates real corruption, not a short/skipped row)"; + } else if (v.size() < expected) { + std::cerr << "Warning: " << name << " breakdown has " << v.size() + << " fields, expected " << expected + << "; padding with -1 to keep the CSV row full width.\n"; + v.resize(expected, -1L); + } + return v; + }; + std::vector cqrrt_bd = fit_breakdown(result.cqrrt.breakdown, 10, "CQRRT_linop"); + std::vector cholqr_bd = fit_breakdown(result.cholqr.breakdown, 5, "CholQR"); + std::vector scholqr3_bd = fit_breakdown(result.scholqr3.breakdown, 17, "sCholQR3"); + std::vector dense_bd = fit_breakdown(result.dense_cqrrt.breakdown, 10, "CQRRT_expl"); + out << std::fixed << std::setprecision(1) << result.m << "," << result.n << "," << run << "," << result.aspect_ratio << "," << std::scientific << std::setprecision(6) << result.cond_num << "," @@ -547,60 +666,107 @@ static void write_results_to_csv( for (const auto* q : algos) out << q->orth_error << "," << q->max_orth_cols << "," << (q->is_orthonormal ? 1 : 0) << "," << q->time << ","; for (int i = 0; i < 4; ++i) - out << algos[i]->peak_rss_kb << "," << algos[i]->analytical_kb << (i < 3 ? "," : "\n"); + out << algos[i]->peak_rss_kb << "," << algos[i]->analytical_kb << ","; + for (int i = 0; i < 4; ++i) + out << algos[i]->qr_status << ","; + for (int i = 0; i < 4; ++i) + out << algos[i]->chol_retries << "," + << algos[i]->chol_shift_abs << "," << algos[i]->chol_shift_rel << (i < 3 ? "," : "\n"); breakdown << result.m << "," << result.n << "," << run << ","; - for (const auto& t : result.cqrrt.breakdown) breakdown << t << ","; + for (const auto& t : cqrrt_bd) breakdown << t << ","; breakdown << result.cqrrt.time << ","; - for (const auto& t : result.cholqr.breakdown) breakdown << t << ","; + for (const auto& t : cholqr_bd) breakdown << t << ","; breakdown << result.cholqr.time << ","; - for (const auto& t : result.scholqr3.breakdown) breakdown << t << ","; + for (const auto& t : scholqr3_bd) breakdown << t << ","; breakdown << result.scholqr3.time << ","; - for (const auto& t : result.dense_cqrrt.breakdown) breakdown << t << ","; + for (const auto& t : dense_bd) breakdown << t << ","; breakdown << result.dense_cqrrt.time << "," << result.cqrrt.peak_rss_kb << "," << result.cqrrt.analytical_kb << "," << result.cholqr.peak_rss_kb << "," << result.cholqr.analytical_kb << "," << result.scholqr3.peak_rss_kb << "," << result.scholqr3.analytical_kb << "," - << result.dense_cqrrt.peak_rss_kb << "," << result.dense_cqrrt.analytical_kb << "\n"; + << result.dense_cqrrt.peak_rss_kb << "," << result.dense_cqrrt.analytical_kb << "," + << result.cqrrt.qr_status << "," << result.cholqr.qr_status << "," + << result.scholqr3.qr_status << "," << result.dense_cqrrt.qr_status << "\n"; } out.flush(); breakdown.flush(); } -// Print console summary for a single size's results +// Print console summary for a single size's results. Only compares runs +// with qr_status == 0: a failed run's `.time` carries the -1 sentinel, which +// must never be picked as "fastest". best_* is -1 when every run +// for that algorithm failed (or, for CQRRT_expl, when skip_dense=1). template static void print_console_summary( const std::vector>& all_runs, int64_t num_runs, int64_t n) { - int64_t best_cqrrt = 0, best_cholqr = 0, best_scholqr3 = 0, best_dense = 0; - for (int64_t r = 1; r < num_runs; ++r) { - if (all_runs[r].cqrrt.time < all_runs[best_cqrrt].cqrrt.time) best_cqrrt = r; - if (all_runs[r].cholqr.time < all_runs[best_cholqr].cholqr.time) best_cholqr = r; - if (all_runs[r].scholqr3.time < all_runs[best_scholqr3].scholqr3.time) best_scholqr3 = r; - if (all_runs[r].dense_cqrrt.time < all_runs[best_dense].dense_cqrrt.time) best_dense = r; + int64_t best_cqrrt = -1, best_cholqr = -1, best_scholqr3 = -1, best_dense = -1; + for (int64_t r = 0; r < num_runs; ++r) { + if (all_runs[r].cqrrt.qr_status == 0 && + (best_cqrrt < 0 || all_runs[r].cqrrt.time < all_runs[best_cqrrt].cqrrt.time)) best_cqrrt = r; + if (all_runs[r].cholqr.qr_status == 0 && + (best_cholqr < 0 || all_runs[r].cholqr.time < all_runs[best_cholqr].cholqr.time)) best_cholqr = r; + if (all_runs[r].scholqr3.qr_status == 0 && + (best_scholqr3 < 0 || all_runs[r].scholqr3.time < all_runs[best_scholqr3].scholqr3.time)) best_scholqr3 = r; + if (all_runs[r].dense_cqrrt.qr_status == 0 && + (best_dense < 0 || all_runs[r].dense_cqrrt.time < all_runs[best_dense].dense_cqrrt.time)) best_dense = r; } - const auto& bc = all_runs[best_cqrrt]; - const auto& bq = all_runs[best_cholqr]; - const auto& bs = all_runs[best_scholqr3]; - const auto& bd = all_runs[best_dense]; - - std::cout << " CQRRT_linop: orth_err=" << std::scientific << std::setprecision(2) << bc.cqrrt.orth_error << ", max_orth=" << bc.cqrrt.max_orth_cols << "/" << n << ", time=" << bc.cqrrt.time << " us (run " << best_cqrrt << ")\n"; - std::cout << " CholQR: orth_err=" << std::scientific << std::setprecision(2) << bq.cholqr.orth_error << ", max_orth=" << bq.cholqr.max_orth_cols << "/" << n << ", time=" << bq.cholqr.time << " us (run " << best_cholqr << ")\n"; - std::cout << " sCholQR3: orth_err=" << std::scientific << std::setprecision(2) << bs.scholqr3.orth_error << ", max_orth=" << bs.scholqr3.max_orth_cols << "/" << n << ", time=" << bs.scholqr3.time << " us (run " << best_scholqr3 << ")\n"; - std::cout << " CQRRT_expl: orth_err=" << std::scientific << std::setprecision(2) << bd.dense_cqrrt.orth_error << ", max_orth=" << bd.dense_cqrrt.max_orth_cols << "/" << n << ", time=" << bd.dense_cqrrt.time << " us (run " << best_dense << ")\n"; + + auto print_alg = [&](const char* label, int64_t best, T orth, int64_t max_orth, long time) { + if (best < 0) { std::cout << " " << label << ": FAILED on every run\n"; return; } + std::cout << " " << label << ": orth_err=" << std::scientific << std::setprecision(2) << orth + << ", max_orth=" << max_orth << "/" << n << ", time=" << time << " us (run " << best << ")\n"; + }; + if (best_cqrrt < 0) print_alg("CQRRT_linop", -1, T(0), 0, 0L); + else print_alg("CQRRT_linop", best_cqrrt, all_runs[best_cqrrt].cqrrt.orth_error, all_runs[best_cqrrt].cqrrt.max_orth_cols, all_runs[best_cqrrt].cqrrt.time); + if (best_cholqr < 0) print_alg("CholQR ", -1, T(0), 0, 0L); + else print_alg("CholQR ", best_cholqr, all_runs[best_cholqr].cholqr.orth_error, all_runs[best_cholqr].cholqr.max_orth_cols, all_runs[best_cholqr].cholqr.time); + if (best_scholqr3 < 0) print_alg("sCholQR3 ", -1, T(0), 0, 0L); + else print_alg("sCholQR3 ", best_scholqr3, all_runs[best_scholqr3].scholqr3.orth_error, all_runs[best_scholqr3].scholqr3.max_orth_cols, all_runs[best_scholqr3].scholqr3.time); + if (best_dense < 0) std::cout << " CQRRT_expl : FAILED or skipped (skip_dense) on every run\n"; + else print_alg("CQRRT_expl ", best_dense, all_runs[best_dense].dense_cqrrt.orth_error, all_runs[best_dense].dense_cqrrt.max_orth_cols, all_runs[best_dense].dense_cqrrt.time); + + // mem_str takes `best` only to decide N/A vs formatted; callers pass 0 for + // rss/akb on the N/A path since those values are never read there (the + // ternary just needs to typecheck without indexing all_runs[-1]). + auto mem_str = [](int64_t best, long rss, long akb) { + return best < 0 ? std::string("N/A") : (std::to_string(rss) + " / " + std::to_string(akb)); + }; + long cqrrt_rss = best_cqrrt < 0 ? 0L : all_runs[best_cqrrt].cqrrt.peak_rss_kb; + long cqrrt_akb = best_cqrrt < 0 ? 0L : all_runs[best_cqrrt].cqrrt.analytical_kb; + long cholqr_rss = best_cholqr < 0 ? 0L : all_runs[best_cholqr].cholqr.peak_rss_kb; + long cholqr_akb = best_cholqr < 0 ? 0L : all_runs[best_cholqr].cholqr.analytical_kb; + long scholqr3_rss = best_scholqr3 < 0 ? 0L : all_runs[best_scholqr3].scholqr3.peak_rss_kb; + long scholqr3_akb = best_scholqr3 < 0 ? 0L : all_runs[best_scholqr3].scholqr3.analytical_kb; + long dense_rss = best_dense < 0 ? 0L : all_runs[best_dense].dense_cqrrt.peak_rss_kb; + long dense_akb = best_dense < 0 ? 0L : all_runs[best_dense].dense_cqrrt.analytical_kb; std::cout << " Memory (peak RSS / analytical KB):\n"; - std::cout << " CQRRT_linop: " << bc.cqrrt.peak_rss_kb << " / " << bc.cqrrt.analytical_kb << ", CholQR: " << bq.cholqr.peak_rss_kb << " / " << bq.cholqr.analytical_kb << ", sCholQR3: " << bs.scholqr3.peak_rss_kb << " / " << bs.scholqr3.analytical_kb << ", CQRRT_expl: " << bd.dense_cqrrt.peak_rss_kb << " / " << bd.dense_cqrrt.analytical_kb << "\n\n"; + std::cout << " CQRRT_linop: " << mem_str(best_cqrrt, cqrrt_rss, cqrrt_akb) + << ", CholQR: " << mem_str(best_cholqr, cholqr_rss, cholqr_akb) + << ", sCholQR3: " << mem_str(best_scholqr3, scholqr3_rss, scholqr3_akb) + << ", CQRRT_expl: " << mem_str(best_dense, dense_rss, dense_akb) + << "\n\n"; } -// Write CSV headers shared by both modes +// Write CSV headers shared by both modes. Provenance lines: argv, +// host, timestamp, and the env knobs that change sCholQR3/CQRRT numerics +// without changing any row label (RANDLAPACK_GRAM_LEFT, RANDLAPACK_SCHOLQR3_ +// SHIFT, RANDLAPACK_GIT_COMMIT among them, via the shared write_env_line). static void write_csv_headers( std::ofstream& out, std::ofstream& breakdown, const std::string& precision, double d_factor, int64_t sketch_nnz, int64_t block_size, int64_t num_runs, int num_threads, - const std::string& extra_comment) { + const std::string& extra_comment, const std::string& argv_line) { + + const std::string run_ts = make_run_timestamp(); + const std::string host = get_hostname(); out << "# CQRRT_linop vs CholQR vs sCholQR3 vs CQRRT_expl Results\n"; + out << "# Date: " << run_ts << "\n"; + out << "# host: " << host << "\n"; + out << "# argv: " << argv_line << "\n"; out << "# Precision: " << precision << "\n"; if (!extra_comment.empty()) out << extra_comment; out << "# d_factor (CQRRT_linop only): " << d_factor << "\n"; @@ -608,6 +774,16 @@ static void write_csv_headers( out << "# block_size (CQRRT_linop, CholQR, sCholQR3): " << block_size << " (0 = full)\n"; out << "# num_runs: " << num_runs << "\n"; out << "# OpenMP threads: " << num_threads << "\n"; + write_env_line(out); + out << "# qr_status: 0 = success; nonzero = the driver's own failure code.\n"; + out << "# On failure the algorithm's orth/max_orth/time/chol_* fields for that row\n"; + out << "# carry the -1 sentinel (never a value-initialized 0, which would misreport\n"; + out << "# as a perfect result). dense_cqrrt_qr_status = -1 also means skip_dense=1\n"; + out << "# (CQRRT_expl was never run, distinct from a real driver failure code).\n"; + out << "# chol_retries/chol_shift_abs/chol_shift_rel: every algorithm here (including\n"; + out << "# dense CQRRT_expl) has an adaptive Cholesky-shift retry mechanism. -1 = the\n"; + out << "# row failed before a shift record existed, or was skipped (dense_cqrrt under\n"; + out << "# skip_dense=1); 0 = the mechanism ran and used an unshifted pass.\n"; out << "# Format: per-run per-algorithm quality metrics (orth_error, max_orth_cols, orth_flag, time), memory (KB)\n"; out << "m,n,run,aspect_ratio,cond_num,density," << "cqrrt_orth_error,cqrrt_max_orth_cols,cqrrt_is_orth,cqrrt_time_us," @@ -617,9 +793,17 @@ static void write_csv_headers( << "cqrrt_peak_rss_kb,cqrrt_analytical_kb," << "cholqr_peak_rss_kb,cholqr_analytical_kb," << "scholqr3_peak_rss_kb,scholqr3_analytical_kb," - << "dense_cqrrt_peak_rss_kb,dense_cqrrt_analytical_kb\n"; + << "dense_cqrrt_peak_rss_kb,dense_cqrrt_analytical_kb," + << "cqrrt_qr_status,cholqr_qr_status,scholqr3_qr_status,dense_cqrrt_qr_status," + << "cqrrt_chol_retries,cqrrt_chol_shift_abs,cqrrt_chol_shift_rel," + << "cholqr_chol_retries,cholqr_chol_shift_abs,cholqr_chol_shift_rel," + << "scholqr3_chol_retries,scholqr3_chol_shift_abs,scholqr3_chol_shift_rel," + << "dense_cqrrt_chol_retries,dense_cqrrt_chol_shift_abs,dense_cqrrt_chol_shift_rel\n"; breakdown << "# Runtime Breakdown for All Algorithms\n"; + breakdown << "# Date: " << run_ts << "\n"; + breakdown << "# host: " << host << "\n"; + breakdown << "# argv: " << argv_line << "\n"; breakdown << "# Precision: " << precision << "\n"; if (!extra_comment.empty()) breakdown << extra_comment; breakdown << "# d_factor (CQRRT_linop only): " << d_factor << "\n"; @@ -627,20 +811,33 @@ static void write_csv_headers( breakdown << "# block_size (CQRRT_linop, CholQR, sCholQR3): " << block_size << " (0 = full)\n"; breakdown << "# num_runs: " << num_runs << "\n"; breakdown << "# OpenMP threads: " << num_threads << "\n"; - breakdown << "# Times are in microseconds\n"; - breakdown << "# CQRRT_linop: alloc, saso, qr, trtri, linop_precond, linop_gram, trmm_gram, potrf, finalize, rest, total\n"; - breakdown << "# CholQR: alloc, materialize, gram, potrf, rest, total\n"; - breakdown << "# sCholQR3: alloc, materialize, gram1, potrf1, trsm1, syrk2, potrf2, update2, syrk3, potrf3, update3, rest, total\n"; + write_env_line(breakdown); + breakdown << "# Times are in microseconds. On a failed or skipped row every field in\n"; + breakdown << "# that algorithm's slot group is -1 (full width preserved);\n"; + breakdown << "# see the per-algorithm *_qr_status column appended at the end.\n"; + // Slot semantics verified against the driver headers (RandLAPACK/drivers/ + // rl_cqrrt.hh, rl_cholqr_linops.hh). The column NAMES below are + // frozen for CSV-name-based-reader compatibility even where a name now + // reads stale: cqrrt_trtri really holds the configured precond_method's + // inversion time (this benchmark always runs TRSM_IDENTITY, never TRTRI); + // cqrrt_linop_precond/cqrrt_linop_gram really hold the Gram-build's fwd/adj + // operator applies; cqrrt_trmm_gram really holds the Gram-forming GEMM + // combine (cholqr_primitive's gemm_dur), not a TRMM; cholqr_materialize/ + // cholqr_gram really hold CholQR's own fwd/adj operator applies. + breakdown << "# CQRRT_linop: alloc, saso, qr, precond_inv, fwd(gram), adj(gram), gemm(gram combine), chol(potrf), finalize, rest, total\n"; + breakdown << "# CholQR: alloc, fwd(gram), adj(gram), chol(potrf), rest, total\n"; + breakdown << "# sCholQR3: alloc, fwd1, adj1, chol1, upd1, fwd2, adj2, gemm2, chol2, upd2, fwd3, adj3, gemm3, chol3, upd3, q_mat, rest, total\n"; breakdown << "# CQRRT_expl: materialize, saso, qr, trtri(=0), precond, gram, trmm_gram(=0), potrf, finalize, rest, total\n"; breakdown << "m,n,run," << "cqrrt_alloc,cqrrt_saso,cqrrt_qr,cqrrt_trtri,cqrrt_linop_precond,cqrrt_linop_gram,cqrrt_trmm_gram,cqrrt_potrf,cqrrt_finalize,cqrrt_rest,cqrrt_total," << "cholqr_alloc,cholqr_materialize,cholqr_gram,cholqr_potrf,cholqr_rest,cholqr_total," - << "scholqr3_alloc,scholqr3_materialize,scholqr3_gram1,scholqr3_potrf1,scholqr3_trsm1,scholqr3_syrk2,scholqr3_potrf2,scholqr3_update2,scholqr3_syrk3,scholqr3_potrf3,scholqr3_update3,scholqr3_rest,scholqr3_total," + << "scholqr3_alloc,scholqr3_fwd1,scholqr3_adj1,scholqr3_chol1,scholqr3_upd1,scholqr3_fwd2,scholqr3_adj2,scholqr3_gemm2,scholqr3_chol2,scholqr3_upd2,scholqr3_fwd3,scholqr3_adj3,scholqr3_gemm3,scholqr3_chol3,scholqr3_upd3,scholqr3_q_mat,scholqr3_rest,scholqr3_total," << "dense_materialize,dense_saso,dense_qr,dense_trtri,dense_precond,dense_gram,dense_trmm_gram,dense_potrf,dense_finalize,dense_rest,dense_total," << "cqrrt_peak_rss_kb,cqrrt_analytical_kb," << "cholqr_peak_rss_kb,cholqr_analytical_kb," << "scholqr3_peak_rss_kb,scholqr3_analytical_kb," - << "dense_cqrrt_peak_rss_kb,dense_cqrrt_analytical_kb\n"; + << "dense_cqrrt_peak_rss_kb,dense_cqrrt_analytical_kb," + << "cqrrt_qr_status,cholqr_qr_status,scholqr3_qr_status,dense_cqrrt_qr_status\n"; } // Prepend total runtime to a CSV file @@ -665,14 +862,23 @@ static int run_benchmark_from_file(int argc, char *argv[]) { // Args: [sketch_nnz] [block_size] [compute_cond] [skip_dense] std::string precision = argv[1]; std::string output_dir = argv[2]; - int64_t num_runs = std::stol(argv[3]); + int64_t num_runs = std::stoll(argv[3]); std::string input_file = argv[4]; double d_factor = std::stod(argv[5]); - int64_t sketch_nnz = (argc >= 7) ? std::stol(argv[6]) : 4; - int64_t block_size = (argc >= 8) ? std::stol(argv[7]) : 0; + int64_t sketch_nnz = (argc >= 7) ? std::stoll(argv[6]) : 4; + int64_t block_size = (argc >= 8) ? std::stoll(argv[7]) : 0; bool compute_cond = (argc >= 9) ? (std::stoi(argv[8]) != 0) : false; bool skip_dense = (argc >= 10) ? (std::stoi(argv[9]) != 0) : false; + // Loud validation: num_runs<=0 leaves `results` empty and + // every per-run write in run_algorithms/write_results_to_csv is UB + // (num_runs=0 was silently accepted before). + if (num_runs < 1) { + std::cerr << "Error: num_runs must be >= 1 (got " << num_runs << ")\n"; + return 1; + } + + std::string argv_line = quote_join_argv(argc, argv); auto benchmark_start = steady_clock::now(); // Generate date/time prefix @@ -715,7 +921,7 @@ static int run_benchmark_from_file(int argc, char *argv[]) { std::string extra_comment = "# Input file: " + input_file + "\n"; write_csv_headers(out, breakdown, precision, d_factor, - sketch_nnz, block_size, num_runs, num_threads, extra_comment); + sketch_nnz, block_size, num_runs, num_threads, extra_comment, argv_line); // Warmup run with a small synthetic matrix { diff --git a/benchmark/bench_CQRRT_linops/CQRRT_linop_composite_applications.cc b/benchmark/bench_CQRRT_linops/CQRRT_linop_composite_applications.cc deleted file mode 100644 index b3cef8c05..000000000 --- a/benchmark/bench_CQRRT_linops/CQRRT_linop_composite_applications.cc +++ /dev/null @@ -1,608 +0,0 @@ -// Generalized SVD / Generalized LS benchmark -// -// Pipeline: -// 1. Load K.mtx (m x m SPD) and V.mtx (m x n sparse) -// 2. Cholesky factorize K = LL^T, create L^{-1} operator (half_solve=true) -// 3. Create composite operator: CompositeOperator(L_inv_op, V_op) = L^{-1}V -// 4. Run Q-less QR on L^{-1}V via CQRRT_linops, CholQR_linops, sCholQR3_linops -// 5. Application (a): Generalized LS — solve min_x ||Vx - b||_{K^{-1}} -// 6. Application (b): Generalized singular values — SVD of R -// 7. Application (c): Generalized singular vectors — full SVD of R -// -// Usage: -// ./GSVD_benchmark -// [sketch_nnz] [block_size] [skip_apps] [compute_cond] - -#include "RandLAPACK.hh" -#include "rl_blaspp.hh" -#include "rl_lapackpp.hh" -#include "rl_gen.hh" - -#include -#include -#include -#include -#include -#include -#include -#include -#ifdef _OPENMP -#include -#endif -#include -#include - -// Extras utilities (Eigen-dependent) -#include "../../extras/misc/ext_util.hh" -#include "../../extras/linops/ext_cholsolver_linop.hh" -#include "RandLAPACK/testing/rl_test_utils.hh" - -// Linops algorithms (now in main RandLAPACK) -#include "rl_cqrrt_linops.hh" -#include "rl_cholqr_linops.hh" -#include "rl_scholqr3_linops.hh" -#include "RandLAPACK/testing/rl_memory_tracker.hh" - -using std::chrono::steady_clock; -using std::chrono::duration_cast; -using std::chrono::microseconds; - -// ============================================================================ -// Result struct -// ============================================================================ - -template -struct gsvd_result { - int64_t m, n; - int64_t run_idx; - std::string alg_name; - - // Cholesky factorization time (shared, measured once) - long chol_time_us; - - // Q-less QR time - long qr_time_us; - - // Orthogonality of Q = (L^{-1}V) R^{-1} - T orth_error; - bool is_orthonormal; - int64_t max_orth_cols; - - // Application (a): Generalized LS - long app_a_time_us; // Post-processing time only - T ls_rel_error; // ||x - x_true|| / ||x_true|| - - // Application (b): Generalized singular values - long app_b_time_us; // SVD of R time - - // Application (c): Generalized singular vectors - long app_c_time_us; // Full SVD of R + V_R orthogonality check - T right_svec_orth_error; // ||V_R^T V_R - I||_F / sqrt(n) - - // Totals (QR + application post-processing) - long total_a_time_us; - long total_b_time_us; - long total_c_time_us; - - // QR timing breakdown (from algo.times[]) - std::vector qr_breakdown; - - // Memory tracking - long peak_rss_kb; // Peak RSS increase during QR call (KB) - long analytical_kb; // Analytical peak working memory (KB) -}; - -// Compute Q = A * R^{-1} uniformly for all algorithms -template -static void compute_Q_from_R( - GLO& A_op, T* R, int64_t ldr, - T* Q_out, int64_t m, int64_t n) { - T* Eye = new T[n * n](); - RandLAPACK::util::eye(n, n, Eye); - A_op(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, - m, n, n, (T)1.0, Eye, n, (T)0.0, Q_out, m); - delete[] Eye; - blas::trsm(blas::Layout::ColMajor, blas::Side::Right, blas::Uplo::Upper, blas::Op::NoTrans, - blas::Diag::NonUnit, m, n, (T)1.0, R, ldr, Q_out, m); -} - -// ============================================================================ -// Application (a): Generalized Least Squares -// min_x ||Vx - b||_{K^{-1}} via R from QR of L^{-1}V -// -// Solution: x = R^{-1} R^{-T} V^T K^{-1} b -// Steps: c = K^{-1}b, d = V^T c, solve R^T y = d, solve Rx = y -// ============================================================================ - -template -static void app_generalized_ls( - RandLAPACK_extras::linops::CholSolverLinOp& K_inv_op, - VLinOp& V_op, - const T* R, int64_t ldr, int64_t n, - const T* b, int64_t m, - T* x, - long& app_time_us) -{ - auto start = steady_clock::now(); - - // Step 1: c = K^{-1} b (m x 1) - std::vector c(m, 0.0); - K_inv_op(blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, - m, 1, m, (T)1.0, b, m, (T)0.0, c.data(), m); - - // Step 2: d = V^T c (n x 1) - std::vector d(n, 0.0); - V_op(blas::Side::Left, blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, - n, 1, m, (T)1.0, c.data(), m, (T)0.0, d.data(), n); - - // Step 3: solve R^T y = d (n x 1) - std::copy(d.begin(), d.end(), x); - blas::trsm(blas::Layout::ColMajor, blas::Side::Left, blas::Uplo::Upper, blas::Op::Trans, - blas::Diag::NonUnit, n, 1, (T)1.0, R, ldr, x, n); - - // Step 4: solve R x = y (n x 1) - blas::trsm(blas::Layout::ColMajor, blas::Side::Left, blas::Uplo::Upper, blas::Op::NoTrans, - blas::Diag::NonUnit, n, 1, (T)1.0, R, ldr, x, n); - - auto stop = steady_clock::now(); - app_time_us = duration_cast(stop - start).count(); -} - -// ============================================================================ -// Application (b): Generalized Singular Values -// SVD of R gives the generalized singular values of (V, K) -// ============================================================================ - -template -static void app_generalized_svals( - const T* R, int64_t ldr, int64_t n, - T* sigma, - long& app_time_us) -{ - auto start = steady_clock::now(); - - // Copy R to work buffer (gesdd destroys input) - std::vector R_copy(n * n); - lapack::lacpy(lapack::MatrixType::General, n, n, R, ldr, R_copy.data(), n); - - // SVD of n x n R: only singular values (no vectors) - std::vector dummy_U(1), dummy_Vt(1); - lapack::gesdd(lapack::Job::NoVec, n, n, R_copy.data(), n, - sigma, dummy_U.data(), 1, dummy_Vt.data(), 1); - - auto stop = steady_clock::now(); - app_time_us = duration_cast(stop - start).count(); -} - -// ============================================================================ -// Application (c): Generalized Singular Vectors -// Full SVD of R: R = U_R * Sigma * V_R^T -// Right generalized singular vectors = columns of V_R -// ============================================================================ - -template -static void app_generalized_svecs( - const T* R, int64_t ldr, int64_t n, - T* sigma, - T* V_R, // n x n right singular vectors - T* U_R, // n x n left singular vectors of R - T& right_svec_orth_error, - long& app_time_us) -{ - auto start = steady_clock::now(); - - // Copy R to work buffer - std::vector R_copy(n * n); - lapack::lacpy(lapack::MatrixType::General, n, n, R, ldr, R_copy.data(), n); - - // Full SVD of n x n R: R = U_R * Sigma * V_R^T - lapack::gesdd(lapack::Job::AllVec, n, n, R_copy.data(), n, - sigma, U_R, n, V_R, n); - - auto stop = steady_clock::now(); - app_time_us = duration_cast(stop - start).count(); - - // Verify right singular vector orthogonality: ||V_R^T V_R - I||_F / sqrt(n) - right_svec_orth_error = RandLAPACK::testing::orthogonality_error(V_R, n, n); -} - -// ============================================================================ -// CSV output -// ============================================================================ - -template -static void write_common_header_comments( - std::ofstream& out, int64_t m, int64_t n, int num_runs, - const std::string& K_file, const std::string& V_file, - T d_factor, int64_t sketch_nnz, int64_t block_size, - bool skip_apps, bool compute_cond) -{ - time_t now = time(nullptr); - out << "# Date: " << ctime(&now) - << "# Matrix dimensions: m=" << m << " n=" << n << "\n" - << "# Runs per algorithm: " << num_runs << "\n" -#ifdef _OPENMP - << "# OpenMP threads: " << omp_get_max_threads() << "\n" -#else - << "# OpenMP threads: 1\n" -#endif - << "# K_file: " << K_file << "\n" - << "# V_file: " << V_file << "\n" - << "# d_factor: " << d_factor << "\n" - << "# sketch_nnz: " << sketch_nnz << "\n" - << "# block_size: " << block_size << "\n" - << "# skip_apps: " << (skip_apps ? 1 : 0) << "\n" - << "# compute_cond: " << (compute_cond ? 1 : 0) << "\n"; -} - -template -static void write_csv_header(std::ofstream& out, int64_t m, int64_t n, int num_runs, - const std::string& K_file, const std::string& V_file, - T d_factor, int64_t sketch_nnz, int64_t block_size, - bool skip_apps, bool compute_cond) { - out << "# GSVD Benchmark results\n"; - write_common_header_comments(out, m, n, num_runs, K_file, V_file, d_factor, sketch_nnz, block_size, skip_apps, compute_cond); - out << "m,n,run,algorithm,chol_time_us,qr_time_us,orth_error,max_orth_cols," - << "app_a_time_us,ls_rel_error," - << "app_b_time_us," - << "app_c_time_us,right_svec_orth_error," - << "total_a_time_us,total_b_time_us,total_c_time_us," - << "peak_rss_kb,analytical_kb\n"; -} - -template -static void write_csv_row(std::ofstream& out, const gsvd_result& r) { - out << r.m << "," << r.n << "," << r.run_idx << "," << r.alg_name << "," - << r.chol_time_us << "," - << r.qr_time_us << "," - << std::scientific << std::setprecision(6) << r.orth_error << "," - << r.max_orth_cols << "," - << r.app_a_time_us << "," - << std::scientific << std::setprecision(6) << r.ls_rel_error << "," - << r.app_b_time_us << "," - << r.app_c_time_us << "," - << std::scientific << std::setprecision(6) << r.right_svec_orth_error << "," - << r.total_a_time_us << "," << r.total_b_time_us << "," << r.total_c_time_us << "," - << r.peak_rss_kb << "," << r.analytical_kb - << "\n"; -} - -// ============================================================================ -// Breakdown CSV output -// ============================================================================ - -template -static void write_breakdown_csv( - const std::string& filename, - const std::vector>& results, - int64_t m, int64_t n, int num_runs, - const std::string& K_file, const std::string& V_file, - T d_factor, int64_t sketch_nnz, int64_t block_size, - bool skip_apps, bool compute_cond) -{ - std::ofstream out(filename); - out << "# GSVD Benchmark runtime breakdown\n"; - write_common_header_comments(out, m, n, num_runs, K_file, V_file, d_factor, sketch_nnz, block_size, skip_apps, compute_cond); - out << "# Times are in microseconds\n"; - out << "# CQRRT_linop breakdown (11): alloc, sketch, qr, tri_inv, fwd, adj, trmm, chol, finalize, rest, total\n"; - out << "# CholQR breakdown (6): alloc, fwd, adj, chol, rest, total\n"; - out << "# sCholQR3 breakdown (18): alloc, fwd1, adj1, chol1, upd1, fwd2, adj2, gemm2, chol2, upd2, fwd3, adj3, gemm3, chol3, upd3, q_mat, rest, total\n"; - out << "# sCholQR3_basic breakdown (15): alloc, fwd1, adj1, chol1, trsm1, fwd_q, syrk2, chol2, upd2, syrk3, chol3, upd3, q_mat, rest, total\n"; - - // Column header: m, n, run, algorithm, then all breakdown times - // Max breakdown length is 18 (sCholQR3) - out << "m,n,run,algorithm"; - for (int i = 0; i < 18; ++i) out << ",t" << i; - out << "\n"; - - for (const auto& r : results) { - out << r.m << "," << r.n << "," << r.run_idx << "," << r.alg_name; - for (size_t i = 0; i < r.qr_breakdown.size(); ++i) { - out << "," << r.qr_breakdown[i]; - } - // Pad with zeros if fewer than 18 columns - for (size_t i = r.qr_breakdown.size(); i < 18; ++i) { - out << ",0"; - } - out << "\n"; - } -} - -// ============================================================================ -// Console summary -// ============================================================================ - -template -static void print_summary(const std::string& alg_name, const std::vector>& results) { - std::cout << "\n " << alg_name.c_str() << ":\n"; - for (const auto& r : results) { - std::cout << " Run " << (long)r.run_idx << ": orth_err=" << std::scientific << std::setprecision(2) << (double)r.orth_error << ", max_orth=" << (long)r.max_orth_cols << "/" << (long)r.n << ", QR=" << r.qr_time_us << " us\n"; - if (r.app_a_time_us > 0 || r.ls_rel_error > 0) { - std::cout << " LS_err=" << std::scientific << std::setprecision(2) << (double)r.ls_rel_error << ", App(a)=" << r.app_a_time_us << " us, App(b)=" << r.app_b_time_us << " us, App(c)=" << r.app_c_time_us << " us\n"; - } - std::cout << " Memory: peak_RSS=" << r.peak_rss_kb << " KB, predicted=" << r.analytical_kb << " KB\n"; - } -} - -// ============================================================================ -// Main benchmark -// ============================================================================ - -template -int run_benchmark(int argc, char* argv[]) { - // Parse arguments - if (argc < 7) { - std::cerr << "Usage: " << argv[0] - << " " - << " [sketch_nnz] [block_size] [skip_apps] [compute_cond]\n"; - return 1; - } - - std::string output_dir = argv[2]; - int64_t num_runs = std::stol(argv[3]); - std::string K_file = argv[4]; - std::string V_file = argv[5]; - T d_factor = std::stod(argv[6]); - int64_t sketch_nnz = (argc >= 8) ? std::stol(argv[7]) : 4; - int64_t block_size = (argc >= 9) ? std::stol(argv[8]) : 0; - bool skip_apps = (argc >= 10) ? (std::stol(argv[9]) != 0) : false; - bool compute_cond = (argc >= 11) ? (std::stol(argv[10]) != 0) : false; - - std::cout << "=== GSVD/Generalized LS Benchmark ===\n"; - std::cout << " K file: " << K_file << "\n"; - std::cout << " V file: " << V_file << "\n"; - std::cout << " d_factor: " << d_factor << "\n"; - std::cout << " sketch_nnz: " << sketch_nnz << "\n"; - std::cout << " block_size: " << block_size << "\n"; - std::cout << " skip_apps: " << (skip_apps ? "yes" : "no") << "\n"; - std::cout << " compute_cond: " << (compute_cond ? "yes" : "no") << "\n"; - std::cout << " num_runs: " << num_runs << "\n"; -#ifdef _OPENMP - std::cout << " OpenMP threads: " << omp_get_max_threads() << "\n\n"; -#else - std::cout << " OpenMP threads: 1\n\n"; -#endif - - // ================================================================ - // Step 1: Load V from Matrix Market - // ================================================================ - std::cout << "Loading V from " << V_file << "... " << std::flush; - auto V_coo = RandLAPACK_extras::coo_from_matrix_market(V_file); - int64_t m = V_coo.n_rows; - int64_t n = V_coo.n_cols; - RandBLAS::sparse_data::csr::CSRMatrix V_csr(m, n); - RandBLAS::sparse_data::conversions::coo_to_csr(V_coo, V_csr); - RandLAPACK::linops::SparseLinOp> V_linop(m, n, V_csr); - std::cout << "done (" << m << " x " << n << ", nnz=" << V_coo.nnz << ")\n"; - - // ================================================================ - // Step 2: Create L^{-1} operator (half_solve=true) and K^{-1} operator - // ================================================================ - std::cout << "Factorizing K = LL^T from " << K_file << "... " << std::flush; - - RandLAPACK_extras::linops::CholSolverLinOp L_inv_op(K_file, /*half_solve=*/true); - auto chol_start = steady_clock::now(); - L_inv_op.factorize(); - auto chol_stop = steady_clock::now(); - long chol_time_us = duration_cast(chol_stop - chol_start).count(); - - // Also create full K^{-1} for App (a) — only needed when running apps - std::unique_ptr> K_inv_op_ptr; - if (!skip_apps) { - K_inv_op_ptr = std::make_unique>(K_file, /*half_solve=*/false); - K_inv_op_ptr->factorize(); - } - - std::cout << "done (" << chol_time_us << " us)\n"; - - // ================================================================ - // Step 3: Form composite operator L^{-1} * V - // ================================================================ - RandLAPACK::linops::CompositeOperator LiV_op(m, n, L_inv_op, V_linop); - LiV_op.block_size = block_size; - std::cout << "Composite operator L^{-1}V: " << m << " x " << n << "\n"; - - // Condition number diagnostic (materializes L^{-1}V, runs two SVDs) - if (compute_cond) { - RandLAPACK::testing::print_condition_diagnostics(LiV_op, "L^{-1}V"); - } - - // ================================================================ - // Step 4: Generate synthetic RHS: b = V * x_true (only when running apps) - // ================================================================ - RandBLAS::RNGState rng_state(42); - std::vector x_true(n); - std::vector b(m, 0.0); - T x_true_norm = 0.0; - if (!skip_apps) { - RandBLAS::DenseDist D(n, 1); - auto next_state = RandBLAS::fill_dense(D, x_true.data(), rng_state); - rng_state = next_state; - - V_linop(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, - m, 1, n, (T)1.0, x_true.data(), n, (T)0.0, b.data(), m); - - x_true_norm = blas::nrm2(n, x_true.data(), 1); - std::cout << "Generated b = V * x_true (||x_true|| = " << x_true_norm << ")\n"; - } - std::cout << "\n"; - - // ================================================================ - // Prepare RNG states for each run - // ================================================================ - RandBLAS::RNGState main_state(123); - std::vector> run_states(num_runs); - for (int64_t r = 0; r < num_runs; ++r) { - run_states[r] = main_state; - if (r > 0) run_states[r].key.incr(r); - } - - // Shared buffers - std::vector Q_buf(m * n); - T tol = std::pow(std::numeric_limits::epsilon(), 0.85); - - // Storage for all results - std::vector> all_results; - - // ================================================================ - // Run warmup (unreported) - // ================================================================ - std::cout << "Running warmup... " << std::flush; - { - auto warmup_state = run_states[0]; - std::vector R_warmup(n * n, 0.0); - RandLAPACK::CQRRT_linops warmup_algo(false, tol, false); - warmup_algo.nnz = sketch_nnz; - warmup_algo.block_size = block_size; - warmup_algo.call(LiV_op, R_warmup.data(), n, d_factor, warmup_state); - } - std::cout << "done\n\n"; - - // ================================================================ - // Common per-algorithm loop: run QR, check orthogonality, run apps - // ================================================================ - // call_algo(res, R, run_idx) fills res.qr_time_us, res.qr_breakdown, - // res.peak_rss_kb, res.analytical_kb — everything else is shared. - auto run_algo = [&](const std::string& name, auto call_algo) { - std::cout << "\n=== " << name << " ===\n"; - std::vector> results(num_runs); - for (int64_t r = 0; r < num_runs; ++r) { - auto& res = results[r]; - res.m = m; res.n = n; res.run_idx = r; - res.alg_name = name; - res.chol_time_us = chol_time_us; - - std::vector R(n * n, 0.0); - call_algo(res, R, r); - - compute_Q_from_R(LiV_op, R.data(), n, Q_buf.data(), m, n); - res.orth_error = RandLAPACK::testing::orthogonality_error(Q_buf.data(), m, n); - res.is_orthonormal = (res.orth_error <= std::pow(std::numeric_limits::epsilon(), (T)0.75)); - res.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_buf.data(), m, n); - - if (!skip_apps) { - std::vector x_computed(n, 0.0); - app_generalized_ls(*K_inv_op_ptr, V_linop, R.data(), n, n, - b.data(), m, x_computed.data(), res.app_a_time_us); - blas::axpy(n, (T)-1.0, x_true.data(), 1, x_computed.data(), 1); - res.ls_rel_error = blas::nrm2(n, x_computed.data(), 1) / x_true_norm; - - std::vector sigma_b(n, 0.0); - app_generalized_svals(R.data(), n, n, sigma_b.data(), res.app_b_time_us); - - std::vector sigma_c(n, 0.0), V_R(n * n, 0.0), U_R(n * n, 0.0); - app_generalized_svecs(R.data(), n, n, sigma_c.data(), V_R.data(), U_R.data(), - res.right_svec_orth_error, res.app_c_time_us); - } - - res.total_a_time_us = res.qr_time_us + res.app_a_time_us; - res.total_b_time_us = res.qr_time_us + res.app_b_time_us; - res.total_c_time_us = res.qr_time_us + res.app_c_time_us; - all_results.push_back(res); - } - print_summary(name, results); - }; - - // ================================================================ - // CQRRT_linop - // ================================================================ - run_algo("CQRRT_linop", [&](gsvd_result& res, std::vector& R, int64_t r) { - auto state = run_states[r]; - RandLAPACK::CQRRT_linops algo(true, tol, false); - algo.nnz = sketch_nnz; algo.block_size = block_size; - RandLAPACK::PeakRSSTracker mem; mem.start(); - algo.call(LiV_op, R.data(), n, d_factor, state); - res.peak_rss_kb = mem.stop(); - res.qr_time_us = algo.times[10]; - // breakdown: alloc, sketch, qr, tri_inv, fwd, adj, trmm, chol, finalize, rest, total - res.qr_breakdown.assign(algo.times.begin(), algo.times.begin() + 11); - res.analytical_kb = RandLAPACK::cqrrt_linops_analytical_kb(m, n, d_factor, block_size); - }); - - // ================================================================ - // CholQR - // ================================================================ - run_algo("CholQR", [&](gsvd_result& res, std::vector& R, int64_t) { - RandLAPACK::CholQR_linops algo(true, tol, false); - algo.block_size = block_size; - RandLAPACK::PeakRSSTracker mem; mem.start(); - algo.call(LiV_op, R.data(), n); - res.peak_rss_kb = mem.stop(); - res.qr_time_us = algo.times[5]; - // breakdown: alloc, fwd, adj, chol, rest, total - res.qr_breakdown.assign(algo.times.begin(), algo.times.begin() + 6); - res.analytical_kb = RandLAPACK::cholqr_linops_analytical_kb(m, n, block_size); - }); - - // ================================================================ - // sCholQR3 - // ================================================================ - run_algo("sCholQR3", [&](gsvd_result& res, std::vector& R, int64_t) { - RandLAPACK::sCholQR3_linops algo(true, tol, false); - algo.block_size = block_size; - RandLAPACK::PeakRSSTracker mem; mem.start(); - algo.call(LiV_op, R.data(), n); - res.peak_rss_kb = mem.stop(); - res.qr_time_us = algo.times[17]; - // breakdown (18): alloc, fwd1, adj1, chol1, upd1, fwd2, adj2, gemm2, chol2, upd2, fwd3, adj3, gemm3, chol3, upd3, q_mat, rest, total - res.qr_breakdown.assign(algo.times.begin(), algo.times.begin() + 18); - res.analytical_kb = RandLAPACK::scholqr3_linops_analytical_kb(m, n, block_size); - }); - - // ================================================================ - // sCholQR3_basic (non-blocked, matches standard pseudocode) - // ================================================================ - run_algo("sCholQR3_basic", [&](gsvd_result& res, std::vector& R, int64_t) { - RandLAPACK::sCholQR3_linops_basic algo(true, tol, false); - RandLAPACK::PeakRSSTracker mem; mem.start(); - algo.call(LiV_op, R.data(), n); - res.peak_rss_kb = mem.stop(); - res.qr_time_us = algo.times[14]; - // breakdown (15): alloc, fwd1, adj1, chol1, trsm1, fwd_q, syrk2, chol2, upd2, syrk3, chol3, upd3, q_mat, rest, total - res.qr_breakdown.assign(algo.times.begin(), algo.times.begin() + 15); - res.analytical_kb = RandLAPACK::scholqr3_linops_basic_analytical_kb(m, n); - }); - - // ================================================================ - // Write CSV output - // ================================================================ - // Generate timestamped filenames - char time_buf[64]; - time_t now = time(nullptr); - strftime(time_buf, sizeof(time_buf), "%Y%m%d_%H%M%S", localtime(&now)); - - std::string results_file = output_dir + "/" + time_buf + "_gsvd_results.csv"; - std::string breakdown_file = output_dir + "/" + time_buf + "_gsvd_breakdown.csv"; - - std::ofstream out(results_file); - write_csv_header(out, m, n, num_runs, K_file, V_file, d_factor, sketch_nnz, block_size, skip_apps, compute_cond); - for (const auto& r : all_results) { - write_csv_row(out, r); - } - out.close(); - std::cout << "\n\nResults written to " << results_file << "\n"; - - write_breakdown_csv(breakdown_file, all_results, m, n, num_runs, K_file, V_file, d_factor, sketch_nnz, block_size, skip_apps, compute_cond); - std::cout << "Runtime breakdown written to " << breakdown_file << "\n"; - - return 0; -} - -int main(int argc, char* argv[]) { - if (argc < 2) { - std::cerr << "Usage: " << argv[0] - << " " - << " [sketch_nnz] [block_size] [skip_apps]\n"; - return 1; - } - - std::string precision = argv[1]; - if (precision == "double") { - return run_benchmark(argc, argv); - } else if (precision == "float") { - return run_benchmark(argc, argv); - } else { - std::cerr << "Unknown precision: " << precision << " (use 'double' or 'float')\n"; - return 1; - } -} diff --git a/benchmark/bench_CQRRT_linops/cqrrt_bench_common.hh b/benchmark/bench_CQRRT_linops/cqrrt_bench_common.hh new file mode 100644 index 000000000..413f38291 --- /dev/null +++ b/benchmark/bench_CQRRT_linops/cqrrt_bench_common.hh @@ -0,0 +1,286 @@ +// cqrrt_bench_common.hh: shared utilities for CQRRT linop benchmarks +#pragma once + +#include "RandLAPACK.hh" +#include "../../extras/misc/ext_util.hh" +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Load a Matrix Market file into a CSRMatrix. Sets m, n, nnz on exit. +template +static RandBLAS::sparse_data::csr::CSRMatrix load_csr( + const std::string& path, int64_t& m, int64_t& n, int64_t& nnz) +{ + auto coo = RandLAPACK_extras::coo_from_matrix_market(path); + m = coo.n_rows; n = coo.n_cols; nnz = coo.nnz; + RandBLAS::sparse_data::csr::CSRMatrix csr(m, n); + RandBLAS::sparse_data::conversions::coo_to_csr(coo, csr); + return csr; +} + +// Load a sparse matrix and emit the standard "Loading