diff --git a/CMakeLists.txt b/CMakeLists.txt index 20687ffc7..67d3db665 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -44,7 +44,10 @@ list(APPEND CMAKE_MODULE_PATH ${CMAKE_SOURCE_DIR}) include(compiler_flags) # Configure the build -enable_testing() +option(BUILD_TESTING "Build tests" ON) +if(BUILD_TESTING) + enable_testing() +endif() include(rl_build_options) include(rl_version) 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..e09ae2ec5 --- /dev/null +++ b/RandLAPACK/comps/rl_cholqr.hh @@ -0,0 +1,558 @@ +#pragma once + +#include "rl_util.hh" +#include "rl_blaspp.hh" +#include "rl_lapackpp.hh" +#include "rl_linops.hh" +#include "rl_bqrrp.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 +}; + + +// ============================================================================ +// 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) 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) { + 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 * m, m); + } + 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 and snapshotted, so retries are O(n^2), not O(m n^2). +// +// 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; nonzero on a diag-zero preconditioner or a Cholesky +// breakdown that survived all retries. +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) +{ + 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); + + 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 (2026-08-12 experiment: 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 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(); } + } + + // ---- 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. We snapshot G and trace(G) once, then on each + // potrf failure restore G, grow the diagonal shift, and retry. trace(G) is + // the natural scale for the shift; G_backup lets 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. + T* G_backup = (max_retries != 0) ? new T[n * n] : nullptr; + T trace_G = 0; + for (int64_t i = 0; i < n; ++i) trace_G += G[i * (n + 1)]; + if (G_backup) std::copy(G, G + n * n, G_backup); + + if (timing) t0 = steady_clock::now(); + + int info = 0; + int attempt = 0; + T current_shift_factor = shift_factor; + // 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 Gram and grow the shift (seed at eps if we started at 0). + std::copy(G_backup, G_backup + n * n, G); + 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; + for (int64_t i = 0; i < n; ++i) G[i * (n + 1)] += shift; + } + if (n > 1) + lapack::laset(MatrixType::Lower, n - 1, n - 1, T(0), T(0), &G[1], n); + info = lapack::potrf(Uplo::Upper, n, G, n); + if (info == 0) break; + } + if (n_retries) *n_retries = attempt; // 0 = succeeded on the unshifted first attempt + + if (info) { + // Report the retries actually made (`attempt`), not the configured limit -- + // with max_retries = -1 the old message printed "-1 retries". + 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[] G_backup; + return info; + } + delete[] G_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) +{ + 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; + + 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); + + 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. +// +// Returns 0 on success, or the 1-based index of the pass whose Cholesky failed. +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 +{ + 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; + 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); + 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); + 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..78e31000e --- /dev/null +++ b/RandLAPACK/drivers/rl_blendenpik.hh @@ -0,0 +1,219 @@ +#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 Oleg asked for (SASO instead of Blendenpik's +// SRFT). It is an INDEPENDENT solver: no mu-regularization, no iterative refinement. +// Reference: knowledge-base/least-squares.md 2.2 "Algorithm SPO1" [BOOK p.101-105]. +// +// Sketch-and-solve initialization (added 2026-07-27, `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_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 (microseconds) + 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; // 0 => default cap (4n) chosen in call() + nnz = 4; // sparse projection, 4 nnz/col (as CQRRT) + lsqr_iters = 0; + warm_start = true; + init_only = false; + } + + /// 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) + { + 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; + auto total_start = clock::now(); + if (init_only) warm_start = true; // x0 is the whole output + + int64_t d = (int64_t)(d_factor * (T)n); + if (d < n) d = n; + + 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 = warm_start ? new T[d]() : nullptr; // sketched RHS, for x0 + T* x0 = warm_start ? new T[n]() : nullptr; + T* r0 = warm_start ? 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) ---- + auto t0 = clock::now(); + 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 (warm_start) { + 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. Cost is O(d n) + O(n^2), negligible + // next to the sketch QR that already happened. + if (warm_start) { + 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]; + } + + // 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}; + } + R_out.assign(R, R + n * n); + cleanup(); + return 0; + } + + // ---- Step 4: LSQR on A with right preconditioner R; x = R^{-1} y ---- + int cap = (max_iters > 0) ? max_iters : (int)std::min(4 * n, 1000); + long lsqr_times[4] = {0, 0, 0, 0}; + t0 = clock::now(); + int st = RandLAPACK::lsqr(A, m, n, R, n, + warm_start ? r0 : b, x, + tol, tol, cap, lsqr_iters, lsqr_times, &final_relres); + if (timing) t_lsqr = duration_cast(clock::now() - t0).count(); + // Undo the shift: LSQR solved for the correction, so add the initial guess back. + if (warm_start) { + 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, + // since b - A(x0 + dx) = r0 - A dx. + if (final_relres >= (T)0) { + T nb = blas::nrm2(m, b, 1); + T nr0 = blas::nrm2(m, r0, 1); + if (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 (previously + // `converged` was written here and never read by any benchmark). + converged = (st == 0); + + if (timing) { + long total = duration_cast(clock::now() - total_start).count(); + this->times = {t_sketch, t_qr, t_lsqr, total}; + } + // Expose the sketch R factor so the caller can report Q = A R^{-1} orthogonality. + R_out.assign(R, R + n * n); + + 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. + std::vector R_out; + /// 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; +}; + + +} // namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_cholqr_dense.hh b/RandLAPACK/drivers/rl_cholqr_dense.hh new file mode 100644 index 000000000..fd7d0a989 --- /dev/null +++ b/RandLAPACK/drivers/rl_cholqr_dense.hh @@ -0,0 +1,287 @@ +#pragma once + +#include "rl_util.hh" +#include "rl_blaspp.hh" +#include "rl_lapackpp.hh" +#include "rl_linops.hh" +#include "../comps/rl_cholqr.hh" + +#include +#include +#include +#include + +using namespace std::chrono; + +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, 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. Ignored when Q is null. +/// @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 +) { + 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)"; + + 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); + if (info != 0) + return info; + + if (Q != nullptr) { + int64_t b_eff = (block_size > 0 && block_size < n) ? block_size : n; + RandLAPACK::materialize_Q_from_R(A_op, R, ldr, m, n, b_eff, Q); + } + 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) + + // 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 = 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 + ) { + steady_clock::time_point total_t_start, total_t_stop; + if (this->timing) total_t_start = steady_clock::now(); + + long it[5] = {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); + if (info != 0) return info; + + if (this->timing) { + total_t_stop = steady_clock::now(); + long total_dur = duration_cast(total_t_stop - total_t_start).count(); + 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; + + // 11 entries: alloc, then [fwd, adj, gemm, chol, upd] per pass, then total + std::vector times; + long total_us() const { return times.empty() ? -1L : times.back(); } + + CholQR2_dense( + bool time_subroutines + ) { + timing = time_subroutines; + block_size = 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 + ) { + steady_clock::time_point total_t_start, total_t_stop; + if (this->timing) total_t_start = steady_clock::now(); + + long it[10] = {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); + if (info != 0) return info; + + if (this->timing) { + total_t_stop = steady_clock::now(); + long total_dur = duration_cast(total_t_stop - total_t_start).count(); + this->times.assign(it, it + 10); + this->times.insert(this->times.begin(), 0L); + this->times.push_back(total_dur); + } + return 0; + } +}; + + +/// Shifted CholeskyQR3 on a dense column-major buffer. +/// Three passes; iter 1 carries the eps shift, iters 2 and 3 are unshifted +/// (Fukaya's prescription). This mirrors sCholQR3_linops exactly. +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; + + // 17 entries: alloc, then [fwd, adj, gemm, chol, upd] per pass, then total + std::vector times; + long total_us() const { return times.empty() ? -1L : times.back(); } + + sCholQR3_dense( + bool time_subroutines + ) { + timing = time_subroutines; + block_size = 0; + shift_factor_iter1 = std::numeric_limits::epsilon(); + 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 + ) { + steady_clock::time_point total_t_start, total_t_stop; + if (this->timing) total_t_start = steady_clock::now(); + + long it[15] = {0}; + int info = detail::cholqr_dense_body( + m, n, A, lda, R, ldr, Q, ldq, + this->block_size, /*num_iters=*/3, + this->shift_factor_iter1, this->shift_factor_iter23, + this->max_retries, this->shift_growth, this->timing, + this->timing ? it : nullptr, &this->n_chol_retries); + if (info != 0) return info; // 1/2/3 = the pass that failed + + if (this->timing) { + total_t_stop = steady_clock::now(); + long total_dur = duration_cast(total_t_stop - total_t_start).count(); + this->times.assign(it, it + 15); + this->times.insert(this->times.begin(), 0L); + this->times.push_back(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..fc7bdd76d 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 { @@ -45,36 +38,26 @@ 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 (Oleg's prescription): 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) + CholQR_linops( bool time_subroutines, T ep, @@ -87,6 +70,8 @@ class CholQR_linops { Q = nullptr; Q_rows = 0; Q_cols = 0; + max_retries = -1; // unbounded retries (no ceiling) + shift_growth = T(10); } ~CholQR_linops() { @@ -95,229 +80,205 @@ 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); + 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); + 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(): +/// 1 if iter-1 cholqr_primitive exhausted retries +/// 2 if iter-2 cholqr_primitive exhausted retries +/// +template +class CholQR2_linops { + public: + bool timing; + bool test_mode; + T eps; - 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) + + CholQR2_linops( + bool time_subroutines, + T ep, + bool enable_test_mode = false + ) { + timing = time_subroutines; + eps = ep; + block_size = 0; + 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); + if (info != 0) return info; // 1 or 2 = the pass that failed + + // ---- 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); + 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; } }; + +// Analytical peak working memory for CholQR2_linops, mirroring the analytic_kb +// helpers used by CholQR_linops / sCholQR3_linops. Sum of class-member scratches: +// G, R_pre, P_prev (3 n^2) + A_temp (m * b_eff) + Z_buf (n * b_eff) +// + G_backup (n^2) when max_retries > 0 +// + cholqr_primitive's transient G + A_temp during iter 1 (n^2 + m*b_eff) +// We report the iter-2 peak (which dominates: persistent class scratches + +// active primitive scratches), conservatively assuming max_retries > 0. +template +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; + int64_t bytes = (int64_t)sizeof(T) * + ( 4 * n * n // G + R_pre + P_prev + G_backup (peak; iter-2 retry) + + 2 * m * b_eff // A_temp (driver) + A_temp (primitive); same allocation in practice + + n * b_eff // Z_buf + ); + return bytes / 1024; +} + } // end namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_cqrrt.hh b/RandLAPACK/drivers/rl_cqrrt.hh index d1dbc55f1..0ff823765 100644 --- a/RandLAPACK/drivers/rl_cqrrt.hh +++ b/RandLAPACK/drivers/rl_cqrrt.hh @@ -6,6 +6,8 @@ #include "rl_lapackpp.hh" #include "rl_hqrrp.hh" #include "rl_bqrrp.hh" +#include "rl_linops.hh" +#include "../comps/rl_cholqr.hh" #include #include @@ -17,6 +19,20 @@ 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: @@ -53,41 +69,9 @@ class CQRRT : public CQRRTalg { /// 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 - /// - int call( int64_t m, int64_t n, @@ -106,16 +90,17 @@ class CQRRT : public CQRRTalg { // 10 entries: saso, qr, trtri(=0), precond, gram, trmm_gram(=0), potrf, finalize, rest, total // Matches CQRRT_linops timing indices for direct comparison. 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) }; // ----------------------------------------------------------------------------- @@ -150,131 +135,75 @@ 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 = d_factor * 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} + 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(); - } - - // Cholesky factorization if (lapack::potrf(Uplo::Upper, n, R_sk, ldr)) { - delete[] A_hat; - delete[] tau; - return 1; + 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 +211,198 @@ 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; + 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 (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; + bool use_dense_sketch; + CQRRTLinopPrecond precond_method; + T bqrrp_block_ratio; + int64_t block_size; + + // Adaptive-shift safety net (Oleg's prescription; 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) + + 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; + 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; + int64_t d = (int64_t)(d_factor * (T)n); + int64_t b_eff = (this->block_size > 0 && this->block_size < n) + ? this->block_size : n; + + // ---- Allocations ---- + 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 = new T[n * n](); + T* G = new T[n * n](); + T* A_temp = new T[m * b_eff]; + T* Z_buf = new T[n * b_eff]; + if (this->timing) { t1 = steady_clock::now(); alloc_dur = duration_cast(t1 - t0).count(); } + + // ---- Step 1: Sketch M^sk = S * A ---- + if (this->timing) t0 = steady_clock::now(); + if (this->use_dense_sketch) { + 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 { + 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 + if (this->timing) { t1 = steady_clock::now(); qr_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); + if (info != 0) { + 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); + 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..e16167acc --- /dev/null +++ b/RandLAPACK/drivers/rl_iter_refine_lsq.hh @@ -0,0 +1,288 @@ +#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. +// +// SINCE 2026-08-07 THIS CLASS IS A THIN ADAPTER over the shared restarted +// engine restarted_pcg_ne (rl_restarted_pcg_ne.hh): the FEM2 and Toeplitz +// benchmarks previously ran two separately-implemented copies of the same +// algorithm (rounds of CG on the right-preconditioned normal equations, +// separated by exact recomputations of the true residual). The 2026-08-06 +// stable-residual change made even the round right-hand sides identical, so +// the outer loops were unified here. What this class adds over the raw engine +// call is the historical field names, the per-step diagnostic vectors, and the +// cold-start policy. +// +// RESTART PACING (Oleg's proposition, 2026-08-07). 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 +// (measured 2026-08-06: inner tolerance is not the accuracy lever, outer +// rounds are). `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 since the 2026-08-06 kernel extraction +// (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 pre-2026-08-07 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 (added 2026-07-29 from ISAAC diagnostic evidence). + /// + /// 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. Set `inner_stag_window <= 0` to disable. + /// + /// WHY, and why this rather than a bigger cap. On the FEM2 operator at + /// kappa^colnorm = 1e10 the CholQR preconditioner is unusable (measured + /// cond(J R^-1) = 7.8e4 against ~1.000 for the other methods, orthogonality error + /// 0.56 against 1e-6). Its inner CG reached its floor at iteration 17 of each + /// 200-iteration step and then ground out the remaining ~183 with no progress. A + /// paired diagnostic run with a 10x larger cap settled the mechanism: + /// + /// cap 200/step (400 total): best_relres 3.483414e-09 @ iter 17, solution error 48.7 + /// cap 2000/step (4000 total): best_relres 3.483414e-09 @ iter 17, solution error 547.0 + /// + /// Ten times the budget left the best residual BIT-IDENTICAL and made the outer + /// solution 11x WORSE, while spending 166 s instead of 13 s. So the cap was never the + /// binding constraint, raising it is actively harmful, and the last iterate is worse + /// than the best one -- hence both halves of this fix. + /// + /// A converging solve is unaffected: it returns Converged before the window elapses. + /// The window is deliberately generous (and the improvement threshold small) so that + /// slow-but-real descent is not mistaken for stagnation. + 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 (Oleg's proposition, 2026-08-07): 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). + /// + /// This replaces the 2026-08-05 `inner_restarts` verification pass: under + /// the restarted scheme every round IS a restart against the true + /// residual, so a separate drift check inside the round is redundant. + 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 (added 2026-08-06, structure unification with + /// restarted_pcg_ne): 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 the check: run exactly + /// n_refine_steps rounds. + T outer_tol; + /// Optional initial iterate to refine (length n), or nullptr for the cold + /// start that every Q-less method uses. Added 2026-08-10 so another solver's + /// answer (Blendenpik's) can be handed to refinement, which separates + /// 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; + /// Final relative residual ||b - J x|| / ||b|| (or ||b - J x|| if ||b|| == 0). + T final_residual_norm; + /// 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. + 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. Always COLD-STARTED: any incoming + /// content is zeroed (the 2026-08-05 policy; the sketch-and-solve + /// warm start is Blendenpik-only). + /// @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 (2026-08-05 policy) unless the caller supplied warm_x0. + std::fill(x, x + n, (T)0); + + // 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); + + // 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; + 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..fed74dbed --- /dev/null +++ b/RandLAPACK/drivers/rl_lsqr.hh @@ -0,0 +1,189 @@ +#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 "rl_lapackpp.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. +/// @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) +{ + using clock = std::chrono::steady_clock; + using std::chrono::duration_cast; + using std::chrono::microseconds; + long t_fwd = 0, t_adj = 0, t_trsm = 0; + auto total_start = clock::now(); + + 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; 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) blas::scal(n, (T)1.0 / alpha, v, 1); + std::copy(v, v + n, w); + + T phibar = beta, rhobar = alpha; + T anorm2 = (T)0; + 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; break; } + if (anorm * rnorm > (T)0 && arnorm <= atol * anorm * rnorm) { status = 0; 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..ad40c8319 --- /dev/null +++ b/RandLAPACK/drivers/rl_pcg_inner.hh @@ -0,0 +1,205 @@ +#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). Extracted 2026-08-06 so the two least-squares drivers +// run the SAME inner solver: stagnation window with best-iterate return (2026-07-29, +// from ISAAC diagnostic evidence), warm-start entry against the true residual +// (2026-08-05), 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). + +#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)) { + rep.iters = it; + rep.status = InnerCGStatus::Breakdown; + rep.relres = std::sqrt(rs_old) / c_norm; + 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..279d4318d --- /dev/null +++ b/RandLAPACK/drivers/rl_restarted_pcg_ne.hh @@ -0,0 +1,403 @@ +#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; Oleg's restart pacing, 2026-08-07 -- +// was 1e-2 until then) and an inner cap of +// min(restart_maxit, max_iters - iters_so_far). Since 2026-08-06 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, or when the inner CG breaks +// down (indefinite H apply, a sign the factor R is unusable). +// +// 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 at round exit + 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 -- Oleg's +/// "CG still terminates once below the target" guard +/// (2026-08-07). 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. +/// Added 2026-08-10 so a solver's own answer (e.g. +/// Blendenpik's) can be handed to iterative refinement, +/// separating preconditioner quality from solver structure. +/// @returns 0 if the LS tolerance was met; 1 if the iteration budget was exhausted; +/// 2 if the inner CG broke down or made no progress (reference flag 2). +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) +{ + 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; + 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](); // true NE residual g - H z + 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 that the + // 2026-08-06 change removed). + T relres; + if (x0 == nullptr) { + std::copy(g, g + n, r_ne); + relres = recover_x_and_relres(); + } else { + std::copy(x0, x0 + n, z); + if (prec) { + auto ts = clock::now(); + 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 + 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); + if (prec) { + Blas2ThreadGuard tg(n); + blas::trsv(blas::Layout::ColMajor, blas::Uplo::Upper, + blas::Op::Trans, blas::Diag::NonUnit, n, R, ldr, r_ne, 1); + } + } + + // 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 (2026-08-07): 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. Two consecutive rounds whose true-residual + // improvement is below stag_rel_improve end the loop with the honest + // not-converged status. Follows the kernel's stagnation switch: disabled + // when 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) break; // round budget spent + T ne_norm = blas::nrm2(n, r_ne, 1); + if (ne_norm == (T)0) { status = 0; break; } // exact NE solution reached + + ++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. + if (inner_abs_tol > (T)0 && ne_norm > (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; + int inner_flag = (kret != 0) ? 2 + : (rep.status == InnerCGStatus::Converged ? 0 : 1); + + 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 (2026-08-06, measured). recover_x_and_relres left + // wm = b - A x (the SMALL true LS residual); map it through R^{-T} A^T + // (Epperly et al. Alg. 1 line 5) instead of the reference's g - H z, which + // subtracts two large kappa-contaminated quantities. On the m=800 prolate + // benchmark case (FFT operator, lambda_rel 1e-20) the two forms were A/B'd + // with everything else identical: g - H z stalls every method's LS relres + // at 1.75e-6 with garbage recovery (1.75e4); this form reaches the 1e-10 + // noise floor with recovery 1.9e-3, matching warm Blendenpik. Dense + // replicas of the same problem do NOT reproduce the stall (both forms + // converge to kappa ~ 1e11 and beyond), so the FFT apply's rounding is a + // necessary ingredient; the benchmark is the regression harness for this. + { + 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); + if (prec) { + { 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); + } + } + } + + 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 (inner_flag == 2) { status = 2; break; } // breakdown: R unusable + if (inner_iters == 0) { status = 2; break; } // no progress possible + if (stag_window > 0) { + if (relres >= ls_stag_ref * ((T)1 - stag_rel_improve)) { + if (++ls_flat_rounds >= 2) break; // LS floor reached: stop + } else { + ls_flat_rounds = 0; + } + if (relres < ls_stag_ref) ls_stag_ref = relres; + } + } + + 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..f2fb67c9b 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,35 @@ #include #include #include +#include +#include + +namespace RandLAPACK { +// Env-gated (read once): RANDLAPACK_SCHOLQR3_SHIFT=theory selects the paper's +// first-pass shift s = 11*eps*n*trace(G) in the sCholQR3 drivers below. +inline bool scholqr3_theory_shift() { + static const bool v = []() { + const char* s = std::getenv("RANDLAPACK_SCHOLQR3_SHIFT"); + return s != nullptr && std::string(s) == "theory"; + }(); + return v; +} +} // namespace RandLAPACK 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. +/// Shifted Cholesky QR3 for abstract linear operators (fully-blocked variant). /// -/// 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 +/// 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 /// -/// 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) -/// -/// 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). /// @@ -59,41 +57,45 @@ class sCholQR3_linops { 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 = eps). Lowered from the original + // `11 * eps * n` to plain `eps` per Oleg: the smaller initial shift avoids + // the iter-2 collapse where shift ≫ σ_min²(A) makes G_2 rank-deficient. // - // 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) sCholQR3_linops( bool time_subroutines, @@ -107,6 +109,10 @@ class sCholQR3_linops { Q = nullptr; Q_rows = 0; Q_cols = 0; + shift_factor_iter1 = std::numeric_limits::epsilon(); // eps*trace(G) shift on iter 1 + 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 +121,82 @@ 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 eps shift + // right away (shift_factor_iter1), iters 2-3 use shift_factor_iter23. + // + // RANDLAPACK_SCHOLQR3_SHIFT=theory switches the first pass to the paper's + // prescription s = 11*eps*n*trace(G) (FukayaEtAl2020, c = 11). History: + // 11*eps*n WAS the original setting, lowered to eps per Oleg after an + // iter-2 collapse (shift >> sigma_min^2 leaves G_2 rank-deficient); this + // knob exists so campaigns can re-measure that trade on real operators. + long it[15] = {0}; + const T sf1 = scholqr3_theory_shift() + ? T(11) * T(n) * std::numeric_limits::epsilon() + : this->shift_factor_iter1; + 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); + if (info != 0) return info; // 1/2/3 = the pass that failed + + // 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); 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 { @@ -469,33 +206,33 @@ class sCholQR3_linops_basic { 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 (Oleg's prescription). + 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) sCholQR3_linops_basic( bool time_subroutines, @@ -508,6 +245,10 @@ class sCholQR3_linops_basic { Q = nullptr; Q_rows = 0; Q_cols = 0; + shift_factor_iter1 = std::numeric_limits::epsilon(); // eps*trace(G) shift on iter 1 + 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 +257,69 @@ 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 RANDLAPACK_SCHOLQR3_SHIFT=theory knob as the blocked variant above. + const T sf1 = scholqr3_theory_shift() + ? T(11) * T(n) * std::numeric_limits::epsilon() + : this->shift_factor_iter1; + 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); + 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); 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..e391e50e2 --- /dev/null +++ b/RandLAPACK/linops/rl_power_linop.hh @@ -0,0 +1,159 @@ +#pragma once + +// Public API: PowerOp — implicit j-th power of a square linear operator. + +#include "rl_concepts.hh" +#include "rl_blaspp.hh" + +#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) + { + randblas_require(base.n_rows == base.n_cols); // PowerOp requires square base + randblas_require(power >= 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) + { + randblas_require(side == Side::Left); + randblas_require(m == n_rows); + randblas_require(k == n_rows); + + if (j == 1) { + base(side, 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(side, 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(side, 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(side, 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 + // for Side::Left (C is N x n) or m x N for Side::Right (C is m x n). + 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) + { + // Determine the shape of the materialized S. + // For Side::Left, op(S) plays the role of B with shape k x n (=> S is k x n or n x k). + // For Side::Right, op(S) is m x k. + 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](); + + // Materialize S via sketch_general(S * I) — works for both DenseSkOp and SparseSkOp. + 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. + (*this)(side, 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..10722339f --- /dev/null +++ b/RandLAPACK/linops/rl_scaled_identity_linop.hh @@ -0,0 +1,92 @@ +#pragma once + +// Public API: ScaledIdentityOp — matrix-free scaled identity mu*I_n. + +#include "rl_exceptions.hh" +#include "rl_concepts.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)"; + + 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..13e2e69bd --- /dev/null +++ b/RandLAPACK/linops/rl_transposed_linop.hh @@ -0,0 +1,99 @@ +#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. + 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; + 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. + 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..8ec2a4c87 --- /dev/null +++ b/RandLAPACK/linops/rl_vstack_linop.hh @@ -0,0 +1,157 @@ +#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; + + 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; + const int64_t bot_off = (layout == Layout::ColMajor) ? r_top : r_top * ldc; + top(side, layout, Op::NoTrans, Op::NoTrans, r_top, n, k, + alpha, B, ldb, beta, C, ldc); + bot(side, 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; + 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(side, layout, Op::Trans, Op::NoTrans, m, n, r_top, + alpha, B, ldb, beta, C, ldc); + bot(side, 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; + + 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..614958a46 --- /dev/null +++ b/RandLAPACK/misc/rl_blas2_threads.hh @@ -0,0 +1,183 @@ +#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 (measured 2026-08-07). +// +// 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 +// +// Three findings pin the diagnosis: +// +// * It is NOT about triangularity, so blocking does not fix it. A plain dgemv of +// the same size degrades the same way (0.234 ms -> 7.04 ms at 16 threads), and a +// hand-blocked solve (sequential diagonal block + threaded gemv update) still +// measured 15.1 ms at 16 threads, ~100x the 4-thread cost. +// * It does NOT improve with size, so there is no crossover to wait for. At 16 +// threads dtrsv costs 46 / 84 / 162 ms for n = 2000 / 4000 / 8000, against +// 0.16 / 1.15 / 4.16 ms at 4 threads. The 16-thread cost grows linearly in n, +// which is the signature of a per-column barrier rather than of arithmetic. +// * At 4 threads all variants sit at the memory-bandwidth floor (~16 MB of the +// factor read per solve), so there is nothing left for an algorithm to win. +// +// CONSEQUENCE IF LEFT UNGUARDED. The preconditioned least-squares solvers apply the +// preconditioner twice per inner iteration; the unpreconditioned baseline applies it +// zero times. The overhead therefore taxes exactly the methods that converge in few +// iterations, and inverts the wall-clock ranking: on a 16000 x 2000 prolate-Toeplitz +// case the CQRRT solve measured 1531 ms against the unpreconditioned 106 ms, and +// re-running the identical binary with OMP_NUM_THREADS=4 moved CQRRT to 10.3 ms +// (149x) while the unpreconditioned baseline barely moved (106 -> 85 ms). +// +// 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, 2026-08-07; 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, 2026-08-10) 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 +/// +/// Two things follow. The mean barely improves past 16 threads, and at 64 the +/// call becomes INTERMITTENTLY unstable: individual DftiComputeForward calls were +/// caught taking 16-32 ms instead of ~0.1 ms (stage-resolved, ~99% of the stall +/// inside the FFT call itself). A solver that converges in a handful of iterations +/// cannot average those stalls out -- measured CQRRT solve 344 ms at 64 threads +/// versus 19 ms at 8, with the unpreconditioned baseline (276 applies) barely +/// affected -- so the artifact scaled INVERSELY with preconditioner quality and +/// inverted the wall-clock ranking. +/// +/// 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 alternative, +/// and is unavailable to us because CG produces one right-hand side at a time). +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; +} + + +/// 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..85f38751a 100644 --- a/RandLAPACK/testing/rl_memory_tracker.hh +++ b/RandLAPACK/testing/rl_memory_tracker.hh @@ -12,6 +12,10 @@ #include #include #include +#include +#if defined(__GLIBC__) || defined(__linux__) +#include // malloc_trim +#endif namespace RandLAPACK { @@ -41,6 +45,20 @@ static inline long get_rss_kb() { class PeakRSSTracker { public: void start() { + // Release freed heap back to the OS before taking the baseline + // (2026-08-05). 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 diagnosed 2026-07-29, and the inflated + // CQRRT_linop storage bars in the Toeplitz figures). 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); @@ -83,38 +101,73 @@ private: // All return memory in KB. // --------------------------------------------------------------------------- -// CQRRT_linops: A_hat(d*n) + tau(n) + R_sk_inv(n*n) + A_pre(m*b_eff) +// CQRRT_linops (TRSM_IDENTITY / GEQP3). +// Peak during the blocked Gram + Cholesky moment: +// A_hat(d*n) + tau(n) + R_sk_inv(n*n) + G(n*n) + G_backup(n*n) + A_pre(m*b_eff) +// = d*n + n + 3*n*n + m*b_eff. The G + G_backup pair is the Cholesky workspace +// and (per the 2026-06-05 rework) the snapshot used for adaptive-shift retries. 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 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); + long bytes = static_cast(sizeof(T)) * (d * n + n + 3L * n * n + (long)m * b_eff); return bytes / 1024; } -// CholQR_linops: I_mat(n*n) + A_temp(m*b_eff) +// CQRRT_linops (BQRRP): the execution has two distinct peak-memory moments: +// (1) BQRRP-preconditioner moment (the column-pivoted-QR inversion inside +// cholqr_primitive, rl_cholqr.hh): A_hat(d*n) + R_sk_inv(n*n) + G(n*n) +// + P_copy(n*n) + R_buf(n*n) = d*n + 4*n*n. A_pre is NOT allocated yet here. +// (W was eliminated by transposing Q^T in place.) +// (2) Gram-loop moment (same as the non-BQRRP path): +// A_hat(d*n) + R_sk_inv(n*n) + tau(n) + A_pre(m*b_eff) +// = d*n + n + n*n + m*b_eff. P_copy/R_buf have been freed by now. +// The true analytical peak is the max of (1) and (2). Roughly: moment (1) wins for +// short-and-wide matrices (m <~ 3n); moment (2) wins for tall matrices. +template +static inline long cqrrt_linops_bqrrp_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 b_eff = (block_size > 0 && block_size < n) ? block_size : n; + long bqrrp_peak = static_cast(sizeof(T)) * ((long)d * n + 4L * n * n); + long gram_peak = static_cast(sizeof(T)) * ((long)d * n + n + (long)n * n + (long)m * b_eff); + return std::max(bqrrp_peak, gram_peak) / 1024; +} + +// CholQR_linops (post-2026-06-05 rework with adaptive-shift retries enabled by default): +// cholqr_primitive owns G(n*n) + G_backup(n*n) + A_temp(m*b_eff) +// plus blocked_preconditioned_gram's I_block(n*b_eff) inside the Gram loop. +// Peak = 2*n*n + (m+n)*b_eff. 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)) * (2L * n * n + (long)(m + n) * b_eff); 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. +// sCholQR3_linops (post-2026-06-05 rework with adaptive-shift retries enabled). +// Persistent driver scratches: G(n*n) + R_pre(n*n) + P_prev(n*n) + A_temp(m*b_eff) +// + Z_buf(n*b_eff) = 3*n*n + (m+n)*b_eff. +// Iter-1 cholqr_primitive also allocates its own G(n*n) + G_backup(n*n) + +// A_temp(m*b_eff) transiently (freed before iters 2/3 start), so the peak is: +// 3*n*n + (m+n)*b_eff + 2*n*n + m*b_eff = 5*n*n + (2m+n)*b_eff. +// Iter-2 / iter-3 cholqr_primitive only adds a single G_backup(n*n), which is +// strictly smaller than iter-1's transient set, so iter-1 dominates the peak. 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)) * (5L * n * n + (long)(2L * m + n) * b_eff); 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 (post-2026-06-05 refactor; non-blocked b_eff = n): +// Same primitive structure as sCholQR3_linops with block_size=0. Plugging +// b_eff = n into the blocked formula 5*n*n + (2m+n)*b_eff gives: +// 5*n*n + (2m + n)*n = 6*n*n + 2*m*n. +// (Identical accounting to the legacy basic formula by coincidence, since the +// b_eff = n collapse cancels the persistent + transient split.) 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)) * (6L * n * n + 2L * (long)m * n); return bytes / 1024; } diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt index 6aa7749d1..8ee8a48b9 100644 --- a/benchmark/CMakeLists.txt +++ b/benchmark/CMakeLists.txt @@ -156,8 +156,44 @@ 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 (Oleg's 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). +add_benchmark(NAME toeplitz_ls_benchmark CXX_SOURCES bench_toeplitz_ls/toeplitz_ls_benchmark.cc LINK_LIBS ${Benchmark_libs_cqrrt_linops}) +# 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(FATAL_ERROR "mkl_dfti.h not found; pass -DMKL_DFTI_INCLUDE_DIR=/include " + "or export MKLROOT for the Toeplitz benchmark.") +endif() +message(STATUS "Toeplitz benchmark: mkl_dfti.h in ${MKL_DFTI_INCLUDE_DIR}") +target_include_directories(toeplitz_ls_benchmark PRIVATE ${MKL_DFTI_INCLUDE_DIR}) +target_compile_definitions(toeplitz_ls_benchmark PRIVATE MKL_ILP64) # 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..57e0f14d4 --- /dev/null +++ b/benchmark/bench_CQRRT_linops/CQRRT_diagnostic.cc @@ -0,0 +1,727 @@ +// 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 — mirrors 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 + TRMM +// 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] + +#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 "../../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; + +// ============================================================================ +// 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 + TRMM + 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)+TRMM + 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{}; + 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); + 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 — see dev log. + 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 = TRMM(R_inv^T, G) [complete Gram: R_inv^T * A^T * A * R_inv] + // ---------------------------------------------------------------- + + // 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, TRMM ---- + // 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) +{ + 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 csv_path = output_dir + "/diagnostic_" + time_buf + ".csv"; + std::ofstream csv(csv_path); + + csv << "# CQRRT Preconditioner Comparison\n"; + csv << "# Date: " << ctime(&now); + 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 << "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 %ld orth_error(Q = A * R_final^{-1}):\n", 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 %ld cond(MR^pre):\n", 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 %ld cond(G = MR^pre^T MR^pre):\n", 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 %ld rel_diff(MR^pre) vs [1]:\n", 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 %ld step-by-step divergence [1] vs [2] (expl: sketch_general; linop: SpGEMM):\n", 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 %ld cond(R_sk): %9.3e\n\n", 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[]) { + 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 output_dir = argv[2]; + 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] + if (argc < 11) { + std::cerr << "Usage (generate mode): " << argv[0] + << " gen [sketch_nnz]\n"; + return 1; + } + int64_t m = std::stol(argv[4]); + int64_t n = std::stol(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::stol(argv[9]); + int64_t sketch_nnz = (argc >= 11) ? std::stol(argv[10]) : 4; + + 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); + } 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::stol(argv[5]); + int64_t sketch_nnz = (argc >= 7) ? std::stol(argv[6]) : 4; + + 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); + } + + 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..60df170e3 --- /dev/null +++ b/benchmark/bench_CQRRT_linops/CQRRT_linop_applications.cc @@ -0,0 +1,2125 @@ +// 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_reg" | "rspec" +// NOTE: these are the ONLY accepted values, and an unrecognized mode is NOT diagnosed -- +// it matches neither dispatch branch, so the run allocates its node, performs no +// experiment, writes no CSV, and exits 0. Six SLURM scripts passing the long-removed +// "irlsq" were archived on 2026-07-29 for exactly this reason. If you add a mode, add it +// here and consider erroring on the unknown case. +// method_mask = bitmask of Q-less QR variants (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 (NOT in the default mask; pass 63 to include it) +// CQRRT_linop_bqrrp (legacy bit 6 / 64) was removed in the 2026-06-05 rework. +// +// Trailing optional args after precond_prec (irlsq / irlsq_reg), added 2026-07-27: +// [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; Oleg's restart +// pacing, 2026-08-07, 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 since 2026-08-07; was 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 (2026-08-05, Max; supersedes the 07-30 2x2 ablation knobs): +// 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 2026-06-09). The former [ir_warm_start] and +// [bp_warm_start] CLI knobs are removed; the 07-30 2x2 that motivated them is +// answered (warm x0 = Blendenpik's forward-error edge, see the 07-30 session log). + +#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 + +// 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 "RandLAPACK/testing/rl_memory_tracker.hh" + +using std::chrono::steady_clock; +using std::chrono::duration_cast; +using std::chrono::microseconds; + +// ============================================================================ +// 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 + int chol_retries = 0; // CholeskyQR adaptive-shift retries (0 = clean; N/A for Blendenpik) + + // 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 INSIDE ir_total_us (0 = cold start) + int ir_outer_iters; + int ir_inner_iters_total; + T ls_residual_norm; + T ls_solution_error; // -1 sentinel when undefined (FEM irlsq) + + // Inner-CG diagnosis (2026-07-27). 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). -1 for Blendenpik (no IR). + // ir_inner_relres : worst per-step achieved ||Mz-c||/||c||. + // 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 when not computed (too large, or not requested). + 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; + + long peak_rss_kb; + long analytical_kb; +}; + +// ---- 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 (2026-07-27): the inner-CG budget was hard-coded at +// 200 per outer step, which with 2 outer steps produced the fixed 400-iteration ceiling +// seen 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 (Oleg 2026-08-07); 0 = legacy fixed-tol rounds +static int g_ir_n_steps = 20; // outer-round cap (2026-08-07, was 4; outer_tol exits early) +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 removed 2026-08-05: IR methods are always +// cold; Blendenpik runs as two mask-32 variants, warm and cold.) + +// 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]; +} + +// Estimate ||A||_2 via power iteration on A^T A. O(iters * (m+n) memory) — no materialization. +template +static T estimate_op_2norm(GLO& A_op, int64_t m, int64_t n, int iters = 10) { + T* v = new T[n]; + T* Av = new T[m]; + { + std::mt19937 rng(7); + std::normal_distribution N01(0, 1); + for (int64_t i = 0; i < n; ++i) v[i] = N01(rng); + } + T sigma = (T)0; + for (int it = 0; it < iters; ++it) { + T nv = blas::nrm2(n, v, 1); + if (nv == 0) { delete[] v; delete[] Av; return (T)0; } + blas::scal(n, (T)1.0 / nv, v, 1); + A_op(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, v, n, (T)0.0, Av, m); + sigma = blas::nrm2(m, Av, 1); + A_op(blas::Side::Left, blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + n, 1, m, (T)1.0, Av, m, (T)0.0, v, n); + } + delete[] v; + delete[] Av; + return sigma; +} + +// orth_err: ||Q^T Q - I||_F / sqrt(n), with Q = A * R^{-1} materialized explicitly. +// O(m*n + n^2) memory. Direct path: materialize Q a column block at a time via the +// linop (so the *peak* extra memory beyond Q itself stays at O(n^2 + m*b)), then +// hand Q off to testing::orthogonality_error. This matches the OLD applications +// benchmark (pre-2026-06-01-b) and the April FEM2 plots' numbers. +// +// The earlier compute_orth_error_memlite path (X = A^T A; X ← R^{-T} X R^{-1}) +// was mathematically equivalent but amplified forward error by kappa(R)^2 from +// the two TRSMs against an ill-conditioned R, producing meaningless ~1e8 values +// on FEM2 where kappa(A) ~ 1e7. The direct path here amplifies only by kappa(R). +// cond_out (optional): when non-null and n <= cond_cap, also returns +// cond(A R^{-1}) = sqrt(lambda_max/lambda_min) of Q^T Q. This is the number that says +// whether the preconditioner actually works, and it costs only one extra n x n syevd on +// top of the Gram this routine already forms -- the expensive part (materializing Q) is +// shared. Added 2026-07-27: App-1 previously logged only max|R_ii|/min|R_ii|, which says +// nothing about cond(A R^{-1}), so a "too many CG iterations" report could not be +// attributed to a weak preconditioner versus an unreachable tolerance. +template +static T compute_orth_error_explicit(GLO& A_op, const T* R, int64_t m, int64_t n, int64_t block_size, + T* cond_out = nullptr, int64_t cond_cap = 16384) { + int64_t b = (block_size > 0 && block_size < n) ? block_size : n; + T* Q = new T[m * n](); + T* E_block = new T[n * b](); // identity column-block scratch + + // Materialize Q = A * R^{-1} one column block at a time: + // E_block = I[:, j:j+b]; Q[:, j:j+b] = A_op * E_block. + // After all blocks, Q = A. Then a single TRSM(Side::Right) gives Q = A * R^{-1}. + for (int64_t j0 = 0; j0 < n; j0 += b) { + int64_t bk = std::min(b, n - j0); + std::fill(E_block, E_block + n * b, (T)0.0); + for (int64_t j = 0; j < bk; ++j) + E_block[(j0 + j) + j * n] = (T)1.0; + A_op(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, bk, n, (T)1.0, E_block, n, (T)0.0, Q + j0 * m, m); + } + delete[] E_block; + + // Q := A * R^{-1} (TRSM amplifies error by kappa(R) only.) + blas::trsm(blas::Layout::ColMajor, blas::Side::Right, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, m, n, (T)1.0, R, n, Q, m); + + T orth = RandLAPACK::testing::orthogonality_error(Q, m, n); + + // cond(A R^{-1}) from the eigenvalues of the Gram, reusing the Q we already built. + // Skipped above cond_cap because syevd is O(n^3) and n reaches 33024 on FEM2-large. + if (cond_out) { + *cond_out = (T)-1; + if (cond_cap <= 0 || n <= cond_cap) { + T* G = new T[n * n](); + blas::syrk(blas::Layout::ColMajor, blas::Uplo::Upper, blas::Op::Trans, + n, m, (T)1.0, Q, m, (T)0.0, G, n); + T* evals = new T[n]; + int info = lapack::syevd(lapack::Job::NoVec, blas::Uplo::Upper, n, G, n, evals); + if (info == 0 && evals[0] > 0) + *cond_out = std::sqrt(evals[n - 1] / evals[0]); // syevd returns ascending + delete[] G; + delete[] evals; + } + } + + delete[] Q; + return orth; +} + +// ============================================================================ +// 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) +{ + std::ofstream out(filename); + time_t now = time(nullptr); + out << "# Sparse IR-LSQ Benchmark results\n" + << "# Date: " << ctime(&now) + << "# input=" << input_label << "\n" + << "# M=" << m << " N=" << n << " nnz=" << nnz_or_zero << "\n" + << "# noise_level=" << noise_level << "\n" + << "# d_factor=" << d_factor << " sketch_nnz=" << sketch_nnz + << " block_size=" << block_size << "\n" + << "# method_mask=" << method_mask << "\n" +#ifdef _OPENMP + << "# OpenMP threads: " << omp_get_max_threads() << "\n" +#else + << "# OpenMP threads: 1\n" +#endif + ; + 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\n"; + 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 + << "\n"; + } +} + +template +static void write_irlsq_breakdown( + const std::string& filename, + const std::vector>& results) +{ + std::ofstream out(filename); + out << "# Sparse IR-LSQ Benchmark runtime breakdown (microseconds)\n" + << "# QR breakdown layout depends on algorithm (see CQRRT_linop_applications.cc).\n" + << "# IR-LSQ breakdown (6): outer_total, inner_cg_total, trsm_total, fwd_total, adj_total, other\n" + << "# (x_0 = 0, so ir_total_us in the results CSV matches outer_total here up to overhead)\n" + << "algorithm,run,phase,t0,t1,t2,t3,t4,t5,t6,t7,t8,t9,t10\n"; + for (const auto& r : results) { + out << r.alg_name << "," << r.run_idx << ",QR"; + for (size_t i = 0; i < r.qr_breakdown.size(); ++i) out << "," << r.qr_breakdown[i]; + for (size_t i = r.qr_breakdown.size(); i < 11; ++i) out << ",0"; + out << "\n"; + out << r.alg_name << "," << r.run_idx << ",IR"; + for (size_t i = 0; i < r.ir_breakdown.size(); ++i) out << "," << r.ir_breakdown[i]; + for (size_t i = r.ir_breakdown.size(); i < 11; ++i) out << ",0"; + out << "\n"; + } +} + +// ============================================================================ +// 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 CQRRT_linop_applications.cc).\n" + << "# RR breakdown (4): orth_error, rayleigh_ritz_build, syevd, ritz_residuals\n" + << "algorithm,run,phase,t0,t1,t2,t3,t4,t5,t6,t7,t8,t9,t10\n"; + for (const auto& r : results) { + out << r.alg_name << "," << r.run_idx << ",QR"; + for (size_t i = 0; i < r.qr_breakdown.size(); ++i) out << "," << r.qr_breakdown[i]; + for (size_t i = r.qr_breakdown.size(); i < 11; ++i) out << ",0"; + out << "\n"; + out << r.alg_name << "," << r.run_idx << ",RR"; + for (size_t i = 0; i < r.rr_breakdown.size(); ++i) out << "," << r.rr_breakdown[i]; + for (size_t i = r.rr_breakdown.size(); i < 11; ++i) out << ",0"; + out << "\n"; + } +} + +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); + time_t now = time(nullptr); + out << "# RSPEC (reduced spectral approximation) results\n" + << "# 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" + << "# 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 %ld (noise=%.3f):\n", + r.alg_name.c_str(), (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: %ld us, peak_RSS=%ld KB, predicted=%ld KB\n", + r.qr_time_us, r.peak_rss_kb, 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=%ld us, outer=%d, inner_total=%d\n", + 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 +{ + // Build the ordered list of selected algorithm names from the bitmask. + // bit 0 ( 1): CQRRT_linop + // bit 1 ( 2): CholQR + // bit 2 ( 4): sCholQR3 + // bit 3 ( 8): sCholQR3_basic + // bit 4 (16): CholQR2 + // bit 5 (32): Blendenpik (sketch + QR + LSQR; independent LSQ solver) + // bit 6 (64): CQRRT_linop_bqrrp [DROPPED in 2026-06-05 rework] + std::vector selected_algs; + if (method_mask & 1) selected_algs.push_back("CQRRT_linop"); + if (method_mask & 2) selected_algs.push_back("CholQR"); + if (method_mask & 4) selected_algs.push_back("sCholQR3"); + if (method_mask & 8) selected_algs.push_back("sCholQR3_basic"); + if (method_mask & 16) selected_algs.push_back("CholQR2"); + if (method_mask & 32) { + selected_algs.push_back("Blendenpik"); // its own sketch-and-solve warm start + selected_algs.push_back("Blendenpik_cold"); // same solver, x0 = 0 + } + if (method_mask & 64) { + // Blendenpik's preconditioner run through OUR refinement solver (2026-08-10). + // As published Blendenpik is sketch + QR + LSQR with no refinement, so the + // suite otherwise compares its R through a different solver than every Q-less + // method uses, conflating preconditioner quality with solver structure. These + // rows keep the published ones intact and add the like-for-like comparison, + // started from Blendenpik's warm and cold answers respectively. + selected_algs.push_back("Blendenpik_refine"); + selected_algs.push_back("Blendenpik_cold_refine"); + } + + 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 (2026-08-05 yellow-bar finding). 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); + if (b_ptr) { + 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||) + T A_2norm = (T)0, b_norm = (T)0; + if (b_ptr) { + std::cout << "Estimating ||A||_2 via power iteration (10 iters)... " << std::flush; + A_2norm = estimate_op_2norm(A_op, m, n, 10); + 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"; + std::vector> alg_results; + + 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 = 0; + + std::fill(R, R + n * n, (T)0); + auto state = run_states[run_idx]; + long blendenpik_lsqr_us = 0; int blendenpik_lsqr_iters = 0; // Blendenpik solve stats + + // ---- 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) { + // Independent sketch-and-precondition LSQ solver: sketch -> QR -> LSQR. + // Produces x_ls directly (no mu, no IR-LSQ). R_out holds the sketch R factor + // used for the Q = A R^{-1} orthogonality check. + RandLAPACK::Blendenpik_linops bp(/*time_subroutines=*/true, tol); + bp.nnz = sketch_nnz; + bp.warm_start = (alg_name == "Blendenpik"); // "_cold" => x0 = 0 + // Match the inner-solve budget the Q-less QR methods get from the IR + // driver (max_inner per step x 2 steps), so the comparison is not + // decided by an accidental cap difference. + bp.max_iters = ((g_ir_max_inner > 0) ? g_ir_max_inner : 200) * g_ir_n_steps; + std::fill(x_ls, x_ls + n, (T)0); + res.qr_status = bp.call(A_op, b_ptr->data(), m, x_ls, n, d_factor, state); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = bp.times[0] + bp.times[1]; // sketch + QR (preconditioner build) + std::copy(bp.R_out.begin(), bp.R_out.end(), R); // sketch R factor for orth + res.qr_breakdown.assign(11, 0); + res.analytical_kb = 0; + // Blendenpik has no IR loop, so reuse the inner-CG diagnosis columns: + // capped = 0/1 from LSQR's own convergence flag, relres = its residual. + res.ir_inner_capped = bp.converged ? 0 : 1; + res.ir_inner_relres = bp.final_relres; + blendenpik_lsqr_us = bp.times[2]; + blendenpik_lsqr_iters = bp.lsqr_iters; + if (alg_name.find("_refine") != std::string::npos) { + // Refine Blendenpik's own answer with the shared engine, so the + // only difference from a Q-less row is which R is used. + T* x_bp = new T[n]; + std::copy(x_ls, x_ls + n, x_bp); + RandLAPACK::IterRefineLSQ ir( + (g_ir_inner_tol > 0) ? (T)g_ir_inner_tol : tol, + (g_ir_max_inner > 0) ? g_ir_max_inner : 200, + g_ir_n_steps, true, 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.warm_x0 = x_bp; + int st = ir.call(A_op, R, n, b_ptr->data(), m, x_ls, n); + if (st != 0) std::cerr << "Warning: refine status " << st << "\n"; + res.ir_outer_iters = ir.outer_iters_done; + record_inner_cg_diagnosis(ir, res); + delete[] x_bp; + } + } + } else if (alg_name == "sCholQR3") { + RandLAPACK::sCholQR3_linops qr_algo(/*time_subroutines=*/true, tol); + qr_algo.block_size = block_size; + res.qr_status = qr_algo.call(A_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.assign(qr_algo.times.begin(), qr_algo.times.begin() + 11); + 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); + res.qr_status = qr_algo.call(A_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.assign(qr_algo.times.begin(), qr_algo.times.begin() + 11); + 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.block_size = block_size; + res.qr_status = qr_algo.call(A_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.assign(qr_algo.times.begin(), qr_algo.times.begin() + 6); + res.qr_breakdown.resize(11, 0); + 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.block_size = block_size; + res.qr_status = qr_algo.call(A_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.assign(qr_algo.times.begin(), qr_algo.times.begin() + 11); + res.analytical_kb = RandLAPACK::cholqr2_linops_analytical_kb(m, n, block_size); + } + } else { + // CQRRT_linop (TRSM_IDENTITY precond). CQRRT_linop_bqrrp was + // removed from the benchmark dispatch in the 2026-06-05 rework. + RandLAPACK::CQRRT_linops qr_algo(/*time_subroutines=*/true, tol); + 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); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = qr_algo.total_us(); + res.qr_breakdown.assign(qr_algo.times.begin(), qr_algo.times.begin() + 11); + 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.qr_breakdown.assign(11, 0); + res.analytical_kb = 0; + alg_results.push_back(res); + 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. ---- + res.orth_error = compute_orth_error_explicit(A_op, R, m, n, block_size, &res.cond_precond); + + // ---- IR-LSQ post-processing ---- + { + const std::vector& b = *b_ptr; + if (alg_name.rfind("Blendenpik", 0) == 0) { + // x_ls was already computed by Blendenpik's own LSQR above; no IR-LSQ. + std::cout << ". LSQR ... " << std::flush; + res.ir_total_us = blendenpik_lsqr_us; + res.ir_outer_iters = 1; + res.ir_inner_iters_total = blendenpik_lsqr_iters; // LSQR iters in the CG-iters slot + } 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(); + 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); + 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); + alg_results.push_back(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"; + write_irlsq_results(results_file, all_results, m, n, input_nnz, input_label, + noise_level, d_factor, sketch_nnz, block_size, method_mask); + std::cout << "\nIR-LSQ results written to " << results_file << "\n"; + write_irlsq_breakdown(breakdown_file, all_results); + std::cout << "IR-LSQ breakdown written to " << breakdown_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) +{ + // Build the list of selected algorithm names from the bitmask (same as run_benchmark_inner). + // bit 6 (64) = CQRRT_linop_bqrrp removed in the 2026-06-05 rework. + std::vector selected_algs; + if (method_mask & 1) selected_algs.push_back("CQRRT_linop"); + if (method_mask & 2) selected_algs.push_back("CholQR"); + if (method_mask & 4) selected_algs.push_back("sCholQR3"); + if (method_mask & 8) selected_algs.push_back("sCholQR3_basic"); + if (method_mask & 16) selected_algs.push_back("CholQR2"); + + 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 = 0; + 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.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.assign(qr_algo.times.begin(), qr_algo.times.begin() + 11); + 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); + 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.assign(qr_algo.times.begin(), qr_algo.times.begin() + 11); + 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.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.assign(qr_algo.times.begin(), qr_algo.times.begin() + 6); + res.qr_breakdown.resize(11, 0); + 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.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.assign(qr_algo.times.begin(), qr_algo.times.begin() + 11); + res.analytical_kb = RandLAPACK::cholqr2_linops_analytical_kb(m, n, block_size); + } + } else { + // CQRRT_linop. CQRRT_linop_bqrrp was removed in 2026-06-05. + RandLAPACK::CQRRT_linops qr_algo(true, tol); + 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.assign(qr_algo.times.begin(), qr_algo.times.begin() + 11); + 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.assign(11, 0); + res.rr_breakdown.assign(4, 0); + 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 — see dev log. + 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(); + res.rspec_total_us = duration_cast(rspec_t1 - rspec_t0).count(); + 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) +{ + std::ofstream out(filename); + time_t now = time(nullptr); + out << "# Sparse IR-LSQ (regularized augmented operator) Benchmark results\n" + << "# Date: " << ctime(&now) + << "# input=" << input_label << "\n" + << "# M=" << m << " N=" << n << " nnz=" << nnz_or_zero << "\n" + << "# d_factor=" << d_factor << " sketch_nnz=" << sketch_nnz + << " block_size=" << block_size << "\n" + << "# method_mask=" << method_mask << "\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; 2026-08-05 policy)\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 + ; + 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\n"; + 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 + << "\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, + 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; + + std::vector selected_algs; + if (method_mask & 1) selected_algs.push_back("CQRRT_linop"); + if (method_mask & 2) selected_algs.push_back("CholQR"); + if (method_mask & 4) selected_algs.push_back("sCholQR3"); + if (method_mask & 8) selected_algs.push_back("sCholQR3_basic"); + if (method_mask & 16) selected_algs.push_back("CholQR2"); + if (method_mask & 32) { // independent sketch+QR+LSQR, warm + cold variants + selected_algs.push_back("Blendenpik"); + selected_algs.push_back("Blendenpik_cold"); + } + if (method_mask & 64) { + // Blendenpik's OWN preconditioner and answer handed to our refinement + // solver (2026-08-10). As published Blendenpik is sketch + QR + LSQR with + // no refinement, so comparing it against Q-less methods that all refine + // conflates preconditioner quality with solver structure. The published + // rows stay untouched; these two add the like-for-like comparison from + // its warm and cold answers. + selected_algs.push_back("Blendenpik_refine"); + selected_algs.push_back("Blendenpik_cold_refine"); + } + 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); + L_inv_Ts.factorize(); + 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); + L_inv_Pp.factorize(); + 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; + + // ||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 before + // 2026-08-05 its thread pools / first-touch pages were paid inside the FIRST + // method's timed solve, which at 4-7 inner iterations is the same magnitude + // as the whole solve (the yellow-bar finding). 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; + + P_precond* R_P = new P_precond[n * n](); + T_solve* R_T = new T_solve[n * n](); + 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; + + 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); + long bp_lsqr_us = 0; int bp_lsqr_iters = 0; + + std::cout << "[Run " << run_idx << ", " << alg_name << "] QR(" << precond_prec_str + << ") ... " << std::flush; + RandLAPACK::PeakRSSTracker mem; mem.start(); + if (is_bp) { + // Independent Blendenpik in SOLVE precision on the BASE operator J_Ts (no mu, + // no augmented A_hat): sketch -> Householder QR -> LSQR, producing x_ls directly. + RandLAPACK::Blendenpik_linops bp(true, tol_T); + bp.nnz = sketch_nnz; + bp.warm_start = (alg_name == "Blendenpik"); // "_cold" => x0 = 0 + // Same inner-solve budget as the IR methods (see the irlsq path). + bp.max_iters = ((g_ir_max_inner > 0) ? g_ir_max_inner : 200) * g_ir_n_steps; + std::fill(x_ls, x_ls + n, (T_solve)0); + res.qr_status = bp.call(J_Ts, b.data(), m, x_ls, n, (T_solve)d_factor, state); + res.peak_rss_kb = mem.stop(); + if (res.qr_status == 0) { + res.qr_time_us = bp.times[0] + bp.times[1]; // sketch + QR (preconditioner build) + std::copy(bp.R_out.begin(), bp.R_out.end(), R_T); // R_T = sketch R factor (orth + kappa) + res.qr_breakdown.assign(11, 0); res.analytical_kb = 0; + res.ir_inner_capped = bp.converged ? 0 : 1; + res.ir_inner_relres = bp.final_relres; + bp_lsqr_us = bp.times[2]; bp_lsqr_iters = bp.lsqr_iters; + if (alg_name.find("_refine") != std::string::npos) { + // Refine Blendenpik's own answer with the shared engine, so the + // only difference from a Q-less row is which R is used. + T_solve* x_bp = new T_solve[n]; + std::copy(x_ls, x_ls + n, x_bp); + 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.warm_x0 = x_bp; + int st_ref = ir.call(J_Ts, R_T, n, b.data(), m, x_ls, n); + if (st_ref != 0) + std::cerr << "Warning: Blendenpik refine status " << st_ref << "\n"; + res.ir_outer_iters = ir.outer_iters_done; + record_inner_cg_diagnosis(ir, res); + bp_lsqr_iters += ir.inner_iters_total(); + delete[] x_bp; + } + } + } else if (alg_name == "sCholQR3") { + RandLAPACK::sCholQR3_linops qr(true, tol_P); qr.block_size = block_size; + res.qr_status = qr.call(A_hat_Pp, R_P, n); res.peak_rss_kb = mem.stop(); res.chol_retries = qr.n_chol_retries; + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown.assign(qr.times.begin(), qr.times.begin() + 11); + res.analytical_kb = RandLAPACK::scholqr3_linops_analytical_kb(m, n, block_size); } + } else if (alg_name == "sCholQR3_basic") { + RandLAPACK::sCholQR3_linops_basic qr(true, tol_P); + res.qr_status = qr.call(A_hat_Pp, R_P, n); res.peak_rss_kb = mem.stop(); res.chol_retries = qr.n_chol_retries; + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown.assign(qr.times.begin(), qr.times.begin() + 11); + res.analytical_kb = RandLAPACK::scholqr3_linops_basic_analytical_kb(m, n); } + } else if (alg_name == "CholQR") { + RandLAPACK::CholQR_linops qr(true, tol_P); qr.block_size = block_size; + res.qr_status = qr.call(A_hat_Pp, R_P, n); res.peak_rss_kb = mem.stop(); res.chol_retries = qr.n_chol_retries; + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown.assign(qr.times.begin(), qr.times.begin() + 6); + res.qr_breakdown.resize(11, 0); + res.analytical_kb = RandLAPACK::cholqr_linops_analytical_kb(m, n, block_size); } + } else if (alg_name == "CholQR2") { + RandLAPACK::CholQR2_linops qr(true, tol_P); qr.block_size = block_size; + res.qr_status = qr.call(A_hat_Pp, R_P, n); res.peak_rss_kb = mem.stop(); res.chol_retries = qr.n_chol_retries; + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown.assign(qr.times.begin(), qr.times.begin() + 11); + res.analytical_kb = RandLAPACK::cholqr2_linops_analytical_kb(m, 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.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.peak_rss_kb = mem.stop(); res.chol_retries = qr.n_chol_retries; + if (res.qr_status == 0) { res.qr_time_us = qr.total_us(); + res.qr_breakdown.assign(qr.times.begin(), qr.times.begin() + 11); + res.analytical_kb = RandLAPACK::cqrrt_linops_analytical_kb(m, 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.qr_breakdown.assign(11, 0); res.analytical_kb = 0; + 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]; + + // Orthogonality loss of Q = A R^{-1} (base A in solve precision). + res.orth_error = compute_orth_error_explicit(J_Ts, R_T, m, n, block_size, &res.cond_precond); + + // Solve in solve precision. Blendenpik already solved (its own LSQR above); + // everyone else runs IR-LSQ with R as the right preconditioner. + if (is_bp) { + std::cout << ". LSQR(" << solve_prec_str << ") ... " << std::flush; + res.ir_total_us = bp_lsqr_us; + res.ir_outer_iters = 1; + res.ir_inner_iters_total = bp_lsqr_iters; // LSQR iters in the CG-iters slot + } else { + std::cout << ". IR-LSQ(" << solve_prec_str << ") ... " << std::flush; + auto ls_t0 = steady_clock::now(); + // Always cold (2026-08-05 policy): the sketch-and-solve x0 warm-start + // ablation that lived here is removed -- warm x0 is Blendenpik-only + // now (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(); + 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); + if (!ir.times.empty()) res.ir_breakdown = ir.times; + } + + // 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"; + 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); + std::cout << "\nIR-LSQ-reg results written to " << results_file << "\n"; + write_irlsq_breakdown(breakdown_file, all_results); + std::cout << "IR-LSQ-reg breakdown written to " << breakdown_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, + 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, + 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, + 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[]) { + // ... → 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::stol(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::stol(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); + int64_t block_size = opt_long(2, 0); + bool compute_cond = (opt_long(3, 0) != 0); + int64_t method_mask = opt_long(4, 31); // bits 0..4: CQRRT_linop, CholQR, sCholQR3, sCholQR3_basic, CholQR2 + 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 (appended 2026-07-27; 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 (Oleg's restart pacing, 2026-08-07). 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] since 2026-08-07 (was " + "[ir_inner_restarts]); it must lie in [0, 1). Regenerate " + "the job scripts.\n"; + return 1; + } + // Outer-round cap (2026-08-07, Max: 20, was 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 instead of being budget-truncated. + g_ir_n_steps = (int)opt_long(14, 20); + if (g_ir_n_steps < 1) { + std::cerr << "Error: ir_n_steps must be >= 1.\n"; + return 1; + } + // Outer early exit (2026-08-06, 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 " + "2026-08-05 (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 + << " (linop=" << (method_mask&1) + << " CholQR=" << ((method_mask>>1)&1) + << " sCholQR3=" << ((method_mask>>2)&1) + << " sCholQR3_basic=" << ((method_mask>>3)&1) + << " CholQR2=" << ((method_mask>>4)&1) << ")\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") { + 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, + 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. + std::cout << "Loading M (mass) from " << M_file << " for X assembly... " << std::flush; + int64_t m_M, n_M, nnz_M; + auto M_csr = load_csr(M_file, m_M, n_M, nnz_M); + std::cout << "done (" << m_M << " x " << n_M << ", nnz=" << nnz_M << ")\n"; + 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.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..5dd88633e 100644 --- a/benchmark/bench_CQRRT_linops/CQRRT_linop_basic.cc +++ b/benchmark/bench_CQRRT_linops/CQRRT_linop_basic.cc @@ -22,9 +22,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" @@ -68,15 +68,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,7 +106,7 @@ 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]; // ============================================================ // Run CQRRT (preconditioned Cholesky QR) - multiple runs @@ -119,36 +117,39 @@ static std::vector> run_algorithms( // RSS measurement (test_mode=false) long cqrrt_peak_rss_kb = 0; { - std::vector R_rss(n * n, 0.0); + T* R_rss = new T[n * n](); 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_rss.call(A_linop, R_rss, n, d_factor, state_rss); cqrrt_peak_rss_kb = cqrrt_mem.stop(); + delete[] R_rss; } + 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); + CQRRT_QR.call(A_linop, R_cqrrt, 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); + 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.data(), m, n); + results[run].cqrrt.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform, m, n); } + delete[] R_cqrrt; } // ============================================================ @@ -160,58 +161,64 @@ static std::vector> run_algorithms( // RSS measurement (test_mode=false) long cholqr_peak_rss_kb = 0; { - std::vector R_rss(n * n, 0.0); + T* R_rss = new T[n * n](); 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_rss.call(A_linop, R_rss, n); cholqr_peak_rss_kb = cholqr_mem.stop(); + delete[] R_rss; } + 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); + CholQR_alg.call(A_linop, R_cholqr, 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); + 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.data(), m, n); + results[run].cholqr.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform, m, 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); + 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); + results[run].scholqr3.time = sCholQR3_alg.times[17]; // total + // 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.begin() + 17); // 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); + 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.data(), m, n); + results[run].scholqr3.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform, m, n); } + delete[] R_scholqr3; } // ============================================================ @@ -220,13 +227,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,17 +243,15 @@ 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); + 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(); @@ -268,11 +274,13 @@ static std::vector> run_algorithms( }; // 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); + 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.data(), m, n); + results[run].dense_cqrrt.max_orth_cols = RandLAPACK::testing::max_orthonormal_cols(Q_uniform, m, n); } + delete[] I_mat; + delete[] R_dense; } // if (!skip_dense) // Compute analytical peak working memory for each algorithm (same for all runs) @@ -287,6 +295,7 @@ static std::vector> run_algorithms( results[r].dense_cqrrt.analytical_kb = dense_cqrrt_akb; } + delete[] Q_uniform; return results; } @@ -354,12 +363,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(); @@ -630,12 +636,12 @@ static void write_csv_headers( 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"; + 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," 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..fc07eb084 --- /dev/null +++ b/benchmark/bench_CQRRT_linops/cqrrt_bench_common.hh @@ -0,0 +1,45 @@ +// 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 + +// 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