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..fc35f1576 100644 --- a/RandLAPACK.hh +++ b/RandLAPACK.hh @@ -28,12 +28,13 @@ #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_iter_refine_lsq.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..4d48f39b5 --- /dev/null +++ b/RandLAPACK/comps/rl_cholqr.hh @@ -0,0 +1,515 @@ +#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 + +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)) +{ + 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: { + // 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 defined(__APPLE__) + // BQRRP cannot link against Apple Accelerate (see rl_bqrrp.hh; the + // sibling rl_cqrrpt.hh / rl_hqrrp.hh guard it the same way). Even + // instantiating RandLAPACK::BQRRP drags in unsupported paths, + // so on macOS fall back to GEQP3, which produces the same column- + // pivoted QR the code below consumes. (void) the BQRRP-only knob. + (void)bqrrp_block_ratio; (void)state; + 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); + #endif + } + + // 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; + } + } + } + + 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. + const bool use_trsm_at_end = preconditioned + && (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; + T current_shift_factor = shift_factor; + for (int attempt = 0; max_retries < 0 || 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(); + } + 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 (info) { + std::fprintf(stderr, + "[cholqr_primitive] FAIL: lapack::potrf returned info=%d after %d " + "retries (final shift_factor=%g)\n", + info, max_retries, (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)) +{ + 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); + + 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) +{ + 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; + 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); + if (info != 0) return 1; + rec(1, fwd1, adj1, 0, chol1, 0); + if (num_iters <= 1) 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); + if (info != 0) { + delete[] G; delete[] R_pre; delete[] P_prev; delete[] A_temp; delete[] Z_buf; + 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; + return 0; +} + +} // namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_cholqr_linops.hh b/RandLAPACK/drivers/rl_cholqr_linops.hh index 5dc7a3494..0962f5daf 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,18 @@ 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. + // 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; + CholQR_linops( bool time_subroutines, T ep, @@ -87,6 +62,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 +72,191 @@ 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); + 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; + 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; - // 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; + + 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); + 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; + 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..319bfe54b 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, @@ -107,15 +91,9 @@ class CQRRT : public CQRRTalg { // Matches CQRRT_linops timing indices for direct comparison. std::vector times; - // 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 +128,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 +204,187 @@ 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; + + 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; + + 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); + 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; + 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..6b3975c1b --- /dev/null +++ b/RandLAPACK/drivers/rl_iter_refine_lsq.hh @@ -0,0 +1,344 @@ +#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, +// and two iterative-refinement steps are performed; under standard hypotheses +// two steps suffice for backward stability. The inner solver is CG on the +// symmetric-positive-definite preconditioned normal-equation matrix +// +// M = R^{-T} J^T J R^{-1}. +// +// 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_lapackpp.hh" +#include "../linops/rl_concepts.hh" + +#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 performing `n_refine_steps` outer steps of +/// +/// r_i ← b - J x_i +/// g_i ← J^T r_i +/// c_i ← R^{-T} g_i +/// z_i ← inner_solve(M, c_i) with M = R^{-T} J^T J R^{-1} +/// x_{i+1} ← x_i + R^{-1} z_i +/// +/// where R is upper triangular (n × n, ColMajor) and is held constant. The +/// inner solver is preconditioner-free conjugate gradients on the SPD matrix +/// M; each inner matvec costs two TRSMs, one J apply (forward), and one J^T +/// apply (adjoint). The default `n_refine_steps = 2` and inner CG stopping +/// rule (relative residual `< inner_tol`) suffice for backward stability +/// under the conditions of Epperly et al. (2025) Theorem 6.1. +template +struct IterRefineLSQ { + // ------------- Configuration ------------- + /// Inner-CG residual tolerance: stop when ||M z - c|| <= inner_tol * ||c||. + T inner_tol; + /// Hard cap on inner CG iterations per outer refinement step. + int max_inner_iters; + /// Outer refinement steps (Algorithm 1 of Epperly et al. uses 2). + int n_refine_steps; + /// Enable per-step / per-substep timing breakdown. + bool timing; + /// Print convergence info to stdout. + bool verbose; + + // ------------- Outputs (filled by call) ------------- + /// Number of outer refinement steps actually executed. + int outer_iters_done; + /// CG iteration counts for each outer step. + std::vector inner_iters_per_step; + /// Final relative residual ||b - J x|| / ||b|| (or ||b - J x|| if ||b|| == 0). + T final_residual_norm; + /// 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), + n_refine_steps(n_steps), + 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. On entry: initial guess (use zeros + /// for cold start). On exit: the refined LS solution. + /// @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) + { + using clock = std::chrono::steady_clock; + auto outer_start = clock::now(); + + // Separate outer-only (excluding inner CG) and inner-CG-only counters so + // we can report a non-overlapping breakdown. The total wallclock spent + // inside inner_cg() is tracked via t_inner_total; the per-op time + // inside inner_cg is captured by the t_inner_* counters and would + // otherwise be triple-counted (once in t_inner_total, once in the + // outer totals) if we shared counters. + long t_inner_total = 0; + long t_outer_trsm = 0, t_outer_fwd = 0, t_outer_adj = 0; + long t_inner_trsm = 0, t_inner_fwd = 0, t_inner_adj = 0; + + inner_iters_per_step.clear(); + outer_iters_done = 0; + + // Per-call workspace buffers. Allocated once up front, reused across outer steps. + T* r = new T[m](); // residual + T* g = new T[n](); // J^T r + T* c = new T[n](); // R^{-T} g + T* z = new T[n](); // inner-solve output + T* dx = new T[n](); // R^{-1} z + T* cg_r = new T[n](); // CG residual + T* cg_p = new T[n](); // CG search direction + T* cg_Mp = new T[n](); // M * p inside CG + T* tmp_n = new T[n](); // R^{-1} p scratch + T* tmp_m = new T[m](); // J * v scratch (m-length) + auto free_workspace = [&]() { + delete[] r; delete[] g; delete[] c; delete[] z; delete[] dx; + delete[] cg_r; delete[] cg_p; delete[] cg_Mp; + delete[] tmp_n; delete[] tmp_m; + }; + + T b_norm = blas::nrm2(m, b, 1); + if (b_norm == (T)0) b_norm = (T)1; // avoid div-by-zero in residual reporting + + for (int step = 0; step < n_refine_steps; ++step) { + // r = b - J*x + // tmp_m = J*x + auto t0 = clock::now(); + J(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, x, n, (T)0.0, tmp_m, m); + t_outer_fwd += std::chrono::duration_cast(clock::now() - t0).count(); + + // r = b - tmp_m + for (int64_t i = 0; i < m; ++i) r[i] = b[i] - tmp_m[i]; + + T r_norm = blas::nrm2(m, r, 1); + if (verbose) { + std::printf("[IR-LSQ] step %d: ||r||/||b|| = %.4e\n", step, (double)(r_norm / b_norm)); + } + + // g = J^T r + t0 = clock::now(); + J(blas::Side::Left, blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + n, 1, m, (T)1.0, r, m, (T)0.0, g, n); + t_outer_adj += std::chrono::duration_cast(clock::now() - t0).count(); + + // c = R^{-T} g (in-place TRSM on a copy of g) + std::copy(g, g + n, c); + t0 = clock::now(); + blas::trsm(blas::Layout::ColMajor, blas::Side::Left, blas::Uplo::Upper, + blas::Op::Trans, blas::Diag::NonUnit, + n, 1, (T)1.0, R, ldr, c, n); + t_outer_trsm += std::chrono::duration_cast(clock::now() - t0).count(); + + // Inner CG on M*z = c + int inner_iters = 0; + auto t_in0 = clock::now(); + int cg_status = inner_cg(J, R, ldr, c, n, m, + z, cg_r, cg_p, cg_Mp, tmp_n, tmp_m, + inner_iters, t_inner_trsm, t_inner_fwd, t_inner_adj); + t_inner_total += std::chrono::duration_cast(clock::now() - t_in0).count(); + inner_iters_per_step.push_back(inner_iters); + if (cg_status != 0) { + outer_iters_done = step; + final_residual_norm = r_norm / b_norm; + if (timing) populate_times(outer_start, t_inner_total, + t_outer_trsm, t_outer_fwd, t_outer_adj, + t_inner_trsm, t_inner_fwd, t_inner_adj); + free_workspace(); + return cg_status; + } + + // dx = R^{-1} z + std::copy(z, z + n, dx); + t0 = clock::now(); + blas::trsm(blas::Layout::ColMajor, blas::Side::Left, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, + n, 1, (T)1.0, R, ldr, dx, n); + t_outer_trsm += std::chrono::duration_cast(clock::now() - t0).count(); + + // x ← x + dx + blas::axpy(n, (T)1.0, dx, 1, x, 1); + outer_iters_done = step + 1; + } + + // Final residual report + { + auto t0 = clock::now(); + J(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, x, n, (T)0.0, tmp_m, m); + t_outer_fwd += std::chrono::duration_cast(clock::now() - t0).count(); + for (int64_t i = 0; i < m; ++i) tmp_m[i] = b[i] - tmp_m[i]; + final_residual_norm = blas::nrm2(m, tmp_m, 1) / b_norm; + } + + if (timing) populate_times(outer_start, t_inner_total, + t_outer_trsm, t_outer_fwd, t_outer_adj, + t_inner_trsm, t_inner_fwd, t_inner_adj); + free_workspace(); + return 0; + } + +private: + // Inner CG: solve M z = c, where M = R^{-T} J^T J R^{-1}, on ℝ^n. + // Workspaces (caller-allocated, length n unless noted): cg_r, cg_p, cg_Mp, + // tmp_n (for R^{-1} v), tmp_m (m-length, for J v_pre). + template + int inner_cg(J_LO& J, const T* R, int64_t ldr, + const T* c, int64_t n, int64_t m, + T* z, + T* cg_r, T* cg_p, T* cg_Mp, + T* tmp_n, T* tmp_m, + int& iters_out, + long& t_trsm, long& t_fwd, long& t_adj) + { + using clock = std::chrono::steady_clock; + // Initial guess z = 0. + std::fill(z, z + n, (T)0); + + // r = c - M*z = c (since z=0) + std::copy(c, c + n, cg_r); + std::copy(c, c + n, cg_p); + + T c_norm = blas::nrm2(n, c, 1); + T tol_abs = inner_tol * c_norm; + if (c_norm == (T)0) { + iters_out = 0; + return 0; + } + + T rs_old = blas::dot(n, cg_r, 1, cg_r, 1); + + for (int it = 0; it < max_inner_iters; ++it) { + // Mp = M * p = R^{-T} J^T J R^{-1} p + // tmp_n = R^{-1} p + std::copy(cg_p, cg_p + n, tmp_n); + auto t0 = clock::now(); + blas::trsm(blas::Layout::ColMajor, blas::Side::Left, blas::Uplo::Upper, + blas::Op::NoTrans, blas::Diag::NonUnit, + n, 1, (T)1.0, R, ldr, tmp_n, n); + t_trsm += std::chrono::duration_cast(clock::now() - t0).count(); + + // tmp_m = J * tmp_n + t0 = clock::now(); + J(blas::Side::Left, blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + m, 1, n, (T)1.0, tmp_n, n, (T)0.0, tmp_m, m); + t_fwd += std::chrono::duration_cast(clock::now() - t0).count(); + + // cg_Mp = J^T * tmp_m + t0 = clock::now(); + J(blas::Side::Left, blas::Layout::ColMajor, blas::Op::Trans, blas::Op::NoTrans, + n, 1, m, (T)1.0, tmp_m, m, (T)0.0, cg_Mp, n); + t_adj += std::chrono::duration_cast(clock::now() - t0).count(); + + // cg_Mp ← R^{-T} cg_Mp (in-place TRSM) + t0 = clock::now(); + blas::trsm(blas::Layout::ColMajor, blas::Side::Left, blas::Uplo::Upper, + blas::Op::Trans, blas::Diag::NonUnit, + n, 1, (T)1.0, R, ldr, cg_Mp, n); + t_trsm += std::chrono::duration_cast(clock::now() - t0).count(); + + T pMp = blas::dot(n, cg_p, 1, cg_Mp, 1); + if (!(pMp > 0)) { + iters_out = it; + return 1; // CG breakdown (loss of orthogonality / non-SPD M) + } + T alpha = rs_old / pMp; + + // z ← z + alpha p + blas::axpy(n, alpha, cg_p, 1, z, 1); + // r ← r - alpha Mp + blas::axpy(n, -alpha, cg_Mp, 1, cg_r, 1); + + T rs_new = blas::dot(n, cg_r, 1, cg_r, 1); + T r_norm = std::sqrt(rs_new); + + if (verbose) { + std::printf("[IR-LSQ] inner CG iter %d: ||r||/||c|| = %.4e\n", + it + 1, (double)(r_norm / c_norm)); + } + if (r_norm <= tol_abs) { + iters_out = it + 1; + return 0; + } + + T beta = rs_new / rs_old; + // p ← r + beta p + for (int64_t i = 0; i < n; ++i) cg_p[i] = cg_r[i] + beta * cg_p[i]; + rs_old = rs_new; + } + iters_out = max_inner_iters; + return 0; // hit cap; not necessarily an error — caller can inspect inner_iters + } + + void populate_times(std::chrono::steady_clock::time_point outer_start, + long t_inner_total, + long t_outer_trsm, long t_outer_fwd, long t_outer_adj, + long t_inner_trsm, long t_inner_fwd, long t_inner_adj) + { + using clock = std::chrono::steady_clock; + long outer_total = std::chrono::duration_cast( + clock::now() - outer_start).count(); + // Non-overlapping breakdown of outer_total: + // inner_ex = inner CG wallclock minus its TRSM/fwd/adj portions + // total_trsm = outer-loop TRSMs + inner-CG TRSMs + // total_fwd = outer-loop J·v + inner-CG J·v + // total_adj = outer-loop J^T·v + inner-CG J^T·v + // other = residual setup, axpy/copy/nrm2 bookkeeping, clock overhead + long inner_ex = t_inner_total - t_inner_trsm - t_inner_fwd - t_inner_adj; + long total_trsm = t_outer_trsm + t_inner_trsm; + long total_fwd = t_outer_fwd + t_inner_fwd; + long total_adj = t_outer_adj + t_inner_adj; + long other = outer_total - inner_ex - total_trsm - total_fwd - total_adj; + times = { outer_total, inner_ex, total_trsm, total_fwd, total_adj, other }; + } +}; + + +} // namespace RandLAPACK diff --git a/RandLAPACK/drivers/rl_scholqr3_linops.hh b/RandLAPACK/drivers/rl_scholqr3_linops.hh index abbef9ad5..e880fafd7 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 @@ -16,33 +17,16 @@ using namespace std::chrono; namespace RandLAPACK { -/// Shifted Cholesky QR3 for abstract linear operators. +/// Shifted Cholesky QR3 for abstract linear operators (fully-blocked variant). /// -/// 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. +/// 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 /// -/// Fully-blocked implementation: never materializes the full m x n operator product -/// during the QR iterations. All three iterations compute their Gram matrices through -/// blocked linop calls, using an accumulated right-factor M = R1^{-1} R2^{-1} ... -/// to avoid storing Q explicitly. -/// -/// Algorithm: -/// 1. Shifted CholQR1: G1 = A^T A + s*I, R1 = chol(G1), M <- R1^{-1} -/// 2. CholQR2: G2 = M^T A^T A M, R2 = chol(G2), R = R2*R1, M <- M*R2^{-1} -/// 3. CholQR3: G3 = M^T A^T A M, R3 = chol(G3), R = R3*R -/// -/// 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 +43,37 @@ 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; - // Column-block size for blocked Gram computations. + int64_t block_size; + + // 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; sCholQR3_linops( bool time_subroutines, @@ -107,6 +87,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 +99,70 @@ 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. + long it[15] = {0}; + int info = cholqr_iterate( + A, R, ldr, 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); + 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; 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,34 +172,26 @@ 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; + // Adaptive shift policy — shared with sCholQR3_linops (Oleg's prescription). + T shift_factor_iter1; + T shift_factor_iter23; + int max_retries; + T shift_growth; + sCholQR3_linops_basic( bool time_subroutines, T ep, @@ -508,6 +203,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 +215,62 @@ 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}; + int info = cholqr_iterate( + A, R, ldr, /*block_size=*/0, /*num_iters=*/3, + this->shift_factor_iter1, this->shift_factor_iter23, + this->max_retries, this->shift_growth, this->timing, + this->timing ? it : nullptr); + 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(); - } + 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/testing/rl_gen.hh b/RandLAPACK/testing/rl_gen.hh index abf6cf812..6ec918b54 100644 --- a/RandLAPACK/testing/rl_gen.hh +++ b/RandLAPACK/testing/rl_gen.hh @@ -514,6 +514,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..af383a1b6 100644 --- a/RandLAPACK/testing/rl_memory_tracker.hh +++ b/RandLAPACK/testing/rl_memory_tracker.hh @@ -12,6 +12,7 @@ #include #include #include +#include namespace RandLAPACK { @@ -83,38 +84,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..2d5be5054 100644 --- a/benchmark/CMakeLists.txt +++ b/benchmark/CMakeLists.txt @@ -156,8 +156,9 @@ 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}) # 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..d586581ea --- /dev/null +++ b/benchmark/bench_CQRRT_linops/CQRRT_diagnostic.cc @@ -0,0 +1,713 @@ +// CQRRT preconditioner comparison benchmark +// +// Isolates the effect of different methods for forming R_sk^{-1} on the final +// orthogonality quality of CQRRT. Tests four paths: +// +// [1] expl_trsm: DTRSM_R(A, R_sk) in-place ← CQRRT_expl path +// [2] expl_inv_trsm: TRSM(I, R_sk) → R_inv; DGEMM(A, R_inv) (TRSM on identity) +// [3] expl_inv_trtri: TRTRI(R_sk) → R_inv; DGEMM(A, R_inv) (LAPACK trtri) +// [4] expl_inv_geqp3: GEQP3(R_sk) = Q*R_buf*P^T; +// R_inv = P * TRSM(R_buf, Q^T); DGEMM(A, R_inv) +// [5] expl_inv_svd: GESDD(R_sk) = U*S*Vt; +// R_inv = V * diag(1/S) * U^T; DGEMM(A, R_inv) +// [6] 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]-[6] all form R_sk^{-1} explicitly via different methods. +// Path [4] 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 [5] uses the SVD (gold standard for stability). +// Path [6] uses BQRRP (blocked randomized QRCP) to invert R_sk — the +// randomized counterpart of path [4]'s GEQP3. +// +// All six 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_12 = ||A_pre[1] - A_pre[2]|| / ||A_pre[1]|| +// rd_13 = ||A_pre[1] - A_pre[3]|| / ||A_pre[1]|| +// rd_14 = ||A_pre[1] - A_pre[4]|| / ||A_pre[1]|| +// rd_15 = ||A_pre[1] - A_pre[5]|| / ||A_pre[1]|| +// rd_16 = ||A_pre[1] - A_pre[6]|| / ||A_pre[1]|| +// +// 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 = 6; + +static constexpr const char* PATH_NAMES[N_PATHS] = { + "expl_trsm", + "expl_inv_trsm", + "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", + "TRSM(I, R_sk)->R_inv; DGEMM(A, R_inv)", + "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: 4-path orth comparison (shared sketch) + +// independent path [1] vs [2] step-by-step divergence +// ============================================================================ + +template +struct TrialResult { + // Per-path metrics (shared sketch, 4 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) + double rd_Apre_12; // TRSM in-place vs TRSM-on-identity + double rd_Apre_13; // TRSM in-place vs trtri + double rd_Apre_14; // TRSM in-place vs geqp3 + double rd_Apre_15; // TRSM in-place vs svd + double rd_Apre_16; // TRSM in-place vs bqrrp + // 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, 4-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 A: TRSM on identity (path [2]) + std::vector R_inv_trsm(n * n, T(0)); + RandLAPACK::util::eye(n, n, R_inv_trsm.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.data(), n); + + // Method B: LAPACK trtri (path [3]) + 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 [4]) + // 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 [5]) + // 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 [6]) + // 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[1]: GEMM + R_inv_trsm (path [2]) + // Apre[2]: GEMM + R_inv_trtri (path [3]) + // Apre[3]: GEMM + R_inv_geqp3 (path [4]) + // ---------------------------------------------------------------- + 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.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]) + // ---------------------------------------------------------------- + res.rd_Apre_12 = (double)rel_diff(Apre[0].data(), Apre[1].data(), m*n); + res.rd_Apre_13 = (double)rel_diff(Apre[0].data(), Apre[2].data(), m*n); + res.rd_Apre_14 = (double)rel_diff(Apre[0].data(), Apre[3].data(), m*n); + res.rd_Apre_15 = (double)rel_diff(Apre[0].data(), Apre[4].data(), m*n); + res.rd_Apre_16 = (double)rel_diff(Apre[0].data(), Apre[5].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 X * R_sk_2 = I (upper triangular result) + std::vector R_inv_2(n * n, T(0)); + RandLAPACK::util::eye(n, n, R_inv_2.data()); + blas::trsm(Layout::ColMajor, Side::Right, 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"; + csv << "run," + << "orth_Q1,orth_Q2,orth_Q3,orth_Q4,orth_Q5,orth_Q6," + << "cond_Apre1,cond_Apre2,cond_Apre3,cond_Apre4,cond_Apre5,cond_Apre6," + << "cond_G1,cond_G2,cond_G3,cond_G4,cond_G5,cond_G6," + << "rd_Apre_12,rd_Apre_13,rd_Apre_14,rd_Apre_15,rd_Apre_16," + << "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); + printf(" rd_12 (trsm-on-I): %12.3e\n", res.rd_Apre_12); + printf(" rd_13 (trtri): %12.3e\n", res.rd_Apre_13); + printf(" rd_14 (geqp3): %12.3e\n", res.rd_Apre_14); + printf(" rd_15 (svd): %12.3e\n", res.rd_Apre_15); + printf(" rd_16 (bqrrp): %12.3e\n", res.rd_Apre_16); + + 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] << ","; + csv << res.rd_Apre_12 << "," << res.rd_Apre_13 << "," << res.rd_Apre_14 << "," + << res.rd_Apre_15 << "," << res.rd_Apre_16 << "," + << 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..762f75956 --- /dev/null +++ b/benchmark/bench_CQRRT_linops/CQRRT_linop_applications.cc @@ -0,0 +1,1741 @@ +// 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" | "rspec" +// 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 +// CQRRT_linop_bqrrp (legacy bit 6 / 64) was removed in the 2026-06-05 rework. + +#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 "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 + + // 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; + int ir_outer_iters; + int ir_inner_iters_total; + T ls_residual_norm; + T ls_solution_error; // -1 sentinel when undefined (FEM irlsq) + + // 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; +}; + +// 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). +template +static T compute_orth_error_explicit(GLO& A_op, const T* R, int64_t m, int64_t n, int64_t block_size) { + 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); + 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\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 + << "\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 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 (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) + 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; + warm_algo.call(A_op, R_warm, n, d_factor, warm_state); + 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]; + + // ---- QR dispatch (5-way, lifted verbatim from CQRRT_linop_irlsq.cc) ---- + std::cout << "[Run " << run_idx << ", " << alg_name << "] QR ... " << std::flush; + RandLAPACK::PeakRSSTracker mem; mem.start(); + 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.times[17]; + 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.times[14]; + 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.times[5]; + 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.times[10]; // total + 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.times[10]; + 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); + + // ---- IR-LSQ post-processing ---- + { + std::cout << ". IR-LSQ ... " << std::flush; + const std::vector& b = *b_ptr; + 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=*/tol, + /*max_inner=*/200, + /*n_steps=*/2, + /*timing=*/true, + /*verbose=*/false); + 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; + 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.times[17]; + 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.times[14]; + 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.times[5]; + 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.times[10]; + 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.times[10]; + 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" + << "# 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\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 + << "\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 (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). + 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; + warm.call(A_hat_Pp, Rw, n, (P_precond)d_factor, ws); delete[] Rw; } + 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]; + + std::cout << "[Run " << run_idx << ", " << alg_name << "] QR(" << precond_prec_str + << ") ... " << std::flush; + RandLAPACK::PeakRSSTracker mem; mem.start(); + 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(); + if (res.qr_status == 0) { res.qr_time_us = qr.times[17]; + 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(); + if (res.qr_status == 0) { res.qr_time_us = qr.times[14]; + 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(); + if (res.qr_status == 0) { res.qr_time_us = qr.times[5]; + 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(); + if (res.qr_status == 0) { res.qr_time_us = qr.times[10]; + 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(); + if (res.qr_status == 0) { res.qr_time_us = qr.times[10]; + 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 = (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. + 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); + + // IR-LSQ in solve precision, R as right preconditioner. + std::cout << ". IR-LSQ(" << solve_prec_str << ") ... " << std::flush; + std::fill(x_ls, x_ls + n, (T_solve)0.0); + auto ls_t0 = steady_clock::now(); + RandLAPACK::IterRefineLSQ ir(tol_T, 200, 2, true, false); + 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; + 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"; + + 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