diff --git a/include/LinearAlgebra/Solvers/tridiagonal_solver.h b/include/LinearAlgebra/Solvers/tridiagonal_solver.h index 5dd47039..c9186b29 100644 --- a/include/LinearAlgebra/Solvers/tridiagonal_solver.h +++ b/include/LinearAlgebra/Solvers/tridiagonal_solver.h @@ -2,6 +2,8 @@ #include +#include + #include "../../LinearAlgebra/Vector/vector.h" #include "../../LinearAlgebra/Vector/vector_operations.h" @@ -9,26 +11,35 @@ namespace gmgpolar { template -class BatchedTridiagonalSolver +class BatchedTridiagonalSolverBase { public: - BatchedTridiagonalSolver(int matrix_dimension, int batch_count, bool is_cyclic = true) + BatchedTridiagonalSolverBase(int matrix_dimension, int batch_count, bool is_cyclic = true) : matrix_dimension_(matrix_dimension) , batch_count_(batch_count) , main_diagonal_("BatchedTridiagonalSolver::main_diagonal", matrix_dimension * batch_count) , sub_diagonal_("BatchedTridiagonalSolver::sub_diagonal", matrix_dimension * batch_count) - , buffer_("BatchedTridiagonalSolver::buffer", is_cyclic ? matrix_dimension * batch_count : 0) - , gamma_("BatchedTridiagonalSolver::gamma", is_cyclic ? batch_count : 0) , is_cyclic_(is_cyclic) - , is_factorized_(false) { + if (matrix_dimension_ <= 0) { + throw std::invalid_argument("matrix_dimension must be positive"); + } + + if (batch_count_ < 0) { + throw std::invalid_argument("batch_count must be non-negative"); + } + assign(main_diagonal_, T(0)); assign(sub_diagonal_, T(0)); } - /* ------------------- */ - /* Accessors for sizes */ - /* ------------------- */ + virtual ~BatchedTridiagonalSolverBase() = default; + + virtual void setup() = 0; + + virtual void solve(Vector rhs, int batch_offset = 0, int batch_stride = 1) = 0; + + virtual void solve_diagonal(Vector rhs, int batch_offset = 0, int batch_stride = 1) = 0; KOKKOS_INLINE_FUNCTION int matrixDimension() const { @@ -40,281 +51,90 @@ class BatchedTridiagonalSolver return batch_count_; } - /* ---------------------------- */ - /* Accessors for matrix entries */ - /* ---------------------------- */ - - KOKKOS_INLINE_FUNCTION const T& main_diagonal(const int batch_idx, const int index) const + KOKKOS_INLINE_FUNCTION + const T& main_diagonal(int batch_idx, int index) const { return main_diagonal_(batch_idx * matrix_dimension_ + index); } - KOKKOS_INLINE_FUNCTION void set_main_diagonal(const int batch_idx, const int index, const T& value) const + + KOKKOS_INLINE_FUNCTION + void set_main_diagonal(int batch_idx, int index, const T& value) const { main_diagonal_(batch_idx * matrix_dimension_ + index) = value; } - KOKKOS_INLINE_FUNCTION void increase_main_diagonal(const int batch_idx, const int index, const T& value) const + + KOKKOS_INLINE_FUNCTION + void increase_main_diagonal(int batch_idx, int index, const T& value) const { main_diagonal_(batch_idx * matrix_dimension_ + index) += value; } - KOKKOS_INLINE_FUNCTION const T& sub_diagonal(const int batch_idx, const int index) const + KOKKOS_INLINE_FUNCTION + const T& sub_diagonal(int batch_idx, int index) const { return sub_diagonal_(batch_idx * matrix_dimension_ + index); } - KOKKOS_INLINE_FUNCTION void set_sub_diagonal(const int batch_idx, const int index, const T& value) const + + KOKKOS_INLINE_FUNCTION + void set_sub_diagonal(int batch_idx, int index, const T& value) const { sub_diagonal_(batch_idx * matrix_dimension_ + index) = value; } - KOKKOS_INLINE_FUNCTION void increase_sub_diagonal(const int batch_idx, const int index, const T& value) const + + KOKKOS_INLINE_FUNCTION + void increase_sub_diagonal(int batch_idx, int index, const T& value) const { sub_diagonal_(batch_idx * matrix_dimension_ + index) += value; } - KOKKOS_INLINE_FUNCTION const T& cyclic_corner(const int batch_idx) const - { - return sub_diagonal_(batch_idx * matrix_dimension_ + (matrix_dimension_ - 1)); - } - KOKKOS_INLINE_FUNCTION T& set_cyclic_corner(const int batch_idx, const T& value) const + KOKKOS_INLINE_FUNCTION + const T& cyclic_corner(int batch_idx) const { - return sub_diagonal_(batch_idx * matrix_dimension_ + (matrix_dimension_ - 1)) = value; + return sub_diagonal_(batch_idx * matrix_dimension_ + matrix_dimension_ - 1); } - KOKKOS_INLINE_FUNCTION void increase_cyclic_corner(const int batch_idx, const T& value) const - { - sub_diagonal_(batch_idx * matrix_dimension_ + (matrix_dimension_ - 1)) += value; - } - - /* ---------------------------------------------- */ - /* Setup: Cholesky Decomposition: A = L * D * L^T */ - /* ---------------------------------------------- */ - // This step factorizes the tridiagonal matrix into lower triangular (L) and diagonal (D) matrices. - // For cyclic systems, it also applies the Shermann-Morrison adjustment to account for the cyclic connection. - void setup() + KOKKOS_INLINE_FUNCTION + void set_cyclic_corner(int batch_idx, const T& value) const { - // Create local copies for lambda capture - int matrix_dimension = matrix_dimension_; - Vector main_diagonal = main_diagonal_; - Vector sub_diagonal = sub_diagonal_; - Vector gamma = gamma_; - - if (!is_cyclic_) { - Kokkos::parallel_for( - "SetupNonCyclic", Kokkos::RangePolicy(0, batch_count_), - KOKKOS_LAMBDA(const int batch_idx) { - // ----------------------------------- // - // Obtain offset for the current batch // - int offset = batch_idx * matrix_dimension; - - // ---------------------- // - // Cholesky Decomposition // - for (int i = 1; i < matrix_dimension; i++) { - sub_diagonal(offset + i - 1) /= main_diagonal(offset + i - 1); - const T factor = sub_diagonal(offset + i - 1); - main_diagonal(offset + i) -= factor * factor * main_diagonal(offset + i - 1); - } - }); - } - else { - Kokkos::parallel_for( - "SetupCyclic", Kokkos::RangePolicy(0, batch_count_), - KOKKOS_LAMBDA(const int batch_idx) { - // ----------------------------------- // - // Obtain offset for the current batch // - int offset = batch_idx * matrix_dimension; - - // ------------------------------------------------- // - // Shermann-Morrison Adjustment // - // - Modify the first and last main diagonal element // - // - Compute and store gamma for later use // - // ------------------------------------------------- // - T cyclic_corner_element = sub_diagonal(offset + matrix_dimension - 1); - gamma(batch_idx) = -main_diagonal(offset + 0); - main_diagonal(offset + 0) -= gamma(batch_idx); - main_diagonal(offset + matrix_dimension - 1) -= - cyclic_corner_element * cyclic_corner_element / gamma(batch_idx); - - // ---------------------- // - // Cholesky Decomposition // - for (int i = 1; i < matrix_dimension; i++) { - sub_diagonal(offset + i - 1) /= main_diagonal(offset + i - 1); - const T factor = sub_diagonal(offset + i - 1); - main_diagonal(offset + i) -= factor * factor * main_diagonal(offset + i - 1); - } - }); - } - Kokkos::fence(); - is_factorized_ = true; + sub_diagonal_(batch_idx * matrix_dimension_ + matrix_dimension_ - 1) = value; } - /* ---------------------------------------- */ - /* Solve: Forward and Backward Substitution */ - /* ---------------------------------------- */ - // This step solves the system Ax = b using the factorized form of A. - // For cyclic systems, it also performs the Shermann-Morrison reconstruction to obtain the final solution. - - void solve(Vector rhs, int batch_offset = 0, int batch_stride = 1) + KOKKOS_INLINE_FUNCTION + void increase_cyclic_corner(int batch_idx, const T& value) const { - if (!is_factorized_) { - throw std::runtime_error("Error: Matrix must be factorized before solving."); - } - - // Compute the effective number of batches to solve - int effective_batch_count = (batch_count_ - batch_offset + batch_stride - 1) / batch_stride; - - // Create local copies for lambda capture - int matrix_dimension = matrix_dimension_; - Vector main_diagonal = main_diagonal_; - Vector sub_diagonal = sub_diagonal_; - Vector buffer = buffer_; - Vector gamma = gamma_; - - if (!is_cyclic_) { - Kokkos::parallel_for( - "SolveNonCyclic", Kokkos::RangePolicy(0, effective_batch_count), - KOKKOS_LAMBDA(const int k) { - // ----------------------------------- // - // Obtain offset for the current batch // - int batch_idx = batch_stride * k + batch_offset; - int offset = batch_idx * matrix_dimension; - - // -------------------- // - // Forward Substitution // - for (int i = 1; i < matrix_dimension; i++) { - rhs(offset + i) -= sub_diagonal(offset + i - 1) * rhs(offset + i - 1); - } - - // ---------------- // - // Diagonal Scaling // - for (int i = 0; i < matrix_dimension; i++) { - rhs(offset + i) /= main_diagonal(offset + i); - } - - // --------------------- // - // Backward Substitution // - for (int i = matrix_dimension - 2; i >= 0; i--) { - rhs(offset + i) -= sub_diagonal(offset + i) * rhs(offset + i + 1); - } - }); - } - else { - Kokkos::parallel_for( - "SolveCyclic", Kokkos::RangePolicy(0, effective_batch_count), - KOKKOS_LAMBDA(const int k) { - // ----------------------------------- // - // Obtain offset for the current batch // - int batch_idx = batch_stride * k + batch_offset; - int offset = batch_idx * matrix_dimension; - - // -------------------- // - // Forward Substitution // - T cyclic_corner_element = sub_diagonal(offset + matrix_dimension - 1); - buffer(offset + 0) = gamma(batch_idx); - for (int i = 1; i < matrix_dimension; i++) { - rhs(offset + i) -= sub_diagonal(offset + i - 1) * rhs(offset + i - 1); - if (i < matrix_dimension - 1) - buffer(offset + i) = 0.0 - sub_diagonal(offset + i - 1) * buffer(offset + i - 1); - else - buffer(offset + i) = - cyclic_corner_element - sub_diagonal(offset + i - 1) * buffer(offset + i - 1); - } - - // ---------------- // - // Diagonal Scaling // - for (int i = 0; i < matrix_dimension; i++) { - rhs(offset + i) /= main_diagonal(offset + i); - buffer(offset + i) /= main_diagonal(offset + i); - } - - // --------------------- // - // Backward Substitution // - for (int i = matrix_dimension - 2; i >= 0; i--) { - rhs(offset + i) -= sub_diagonal(offset + i) * rhs(offset + i + 1); - buffer(offset + i) -= sub_diagonal(offset + i) * buffer(offset + i + 1); - } - - // ------------------------------- // - // Shermann-Morrison Reonstruction // - const T dot_product_x_v = - rhs(offset + 0) + cyclic_corner_element / gamma(batch_idx) * rhs(offset + matrix_dimension - 1); - const T dot_product_u_v = buffer(offset + 0) + cyclic_corner_element / gamma(batch_idx) * - buffer(offset + matrix_dimension - 1); - const T factor = dot_product_x_v / (1.0 + dot_product_u_v); - - for (int i = 0; i < matrix_dimension; i++) { - rhs(offset + i) -= factor * buffer(offset + i); - } - }); - } - Kokkos::fence(); - } - - /* ---------------------------- */ - /* Solve: Diagonal Scaling Only */ - /* ---------------------------- */ - // This step performs only the diagonal scaling part of the solve process. - // It is useful when the matrix has a non-zero diagonal but zero off-diagonal entries. - // Note that .setup() modifies main_diagonal(0) in the cyclic case. - - void solve_diagonal(Vector rhs, int batch_offset = 0, int batch_stride = 1) - { - if (!is_factorized_) { - throw std::runtime_error("Error: Matrix must be factorized before solving."); - } - - // Compute the effective number of batches to solve - int effective_batch_count = (batch_count_ - batch_offset + batch_stride - 1) / batch_stride; - - // Create local copies for lambda capture - int matrix_dimension = matrix_dimension_; - Vector main_diagonal = main_diagonal_; - Vector gamma = gamma_; - - if (!is_cyclic_) { - Kokkos::parallel_for( - "SolveDiagonalNonCyclic", Kokkos::RangePolicy(0, effective_batch_count), - KOKKOS_LAMBDA(const int k) { - // ----------------------------------- // - // Obtain offset for the current batch // - int batch_idx = batch_stride * k + batch_offset; - int offset = batch_idx * matrix_dimension; - - // ---------------- // - // Diagonal Scaling // - for (int i = 0; i < matrix_dimension; i++) { - rhs(offset + i) /= main_diagonal(offset + i); - } - }); - } - else { - Kokkos::parallel_for( - "SolveDiagonalCyclic", Kokkos::RangePolicy(0, effective_batch_count), - KOKKOS_LAMBDA(const int k) { - // ----------------------------------- // - // Obtain offset for the current batch // - int batch_idx = batch_stride * k + batch_offset; - int offset = batch_idx * matrix_dimension; - - // ---------------- // - // Diagonal Scaling // - rhs(offset + 0) /= main_diagonal(offset + 0) + gamma(batch_idx); - for (int i = 1; i < matrix_dimension; i++) { - rhs(offset + i) /= main_diagonal(offset + i); - } - }); - } - Kokkos::fence(); + sub_diagonal_(batch_idx * matrix_dimension_ + matrix_dimension_ - 1) += value; } -private: +protected: int matrix_dimension_; int batch_count_; Vector main_diagonal_; Vector sub_diagonal_; - Vector buffer_; - Vector gamma_; bool is_cyclic_; - bool is_factorized_; }; + +} // namespace gmgpolar + +#include "tridiagonal_solver_thomas.h" +#include "tridiagonal_solver_pcr.h" +#include "tridiagonal_solver_cr.h" +#include "tridiagonal_solver_crpcr.h" + +namespace gmgpolar +{ + +#if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || defined(KOKKOS_ENABLE_SYCL) + +template +using BatchedTridiagonalSolver = BatchedTridiagonalSolverCRPCR; + +#else + +template +using BatchedTridiagonalSolver = BatchedTridiagonalSolverThomas; + +#endif + } // namespace gmgpolar diff --git a/include/LinearAlgebra/Solvers/tridiagonal_solver_cr.h b/include/LinearAlgebra/Solvers/tridiagonal_solver_cr.h new file mode 100644 index 00000000..458f07c8 --- /dev/null +++ b/include/LinearAlgebra/Solvers/tridiagonal_solver_cr.h @@ -0,0 +1,544 @@ +#pragma once + +/** + * @brief Batched tridiagonal solver based on classical Cyclic Reduction (CR). + * + * Classical CR reduces the active system size by approximately a factor of + * two per stage. Unlike Parallel Cyclic Reduction (PCR), CR does not update + * every equation at every stage: at each level, exactly half of the + * currently active equations are eliminated and folded into the equations + * that survive to the next, coarser level. Reduction continues until a + * single root equation remains. + * + * The solver supports arbitrary matrix dimensions without power-of-two + * padding. Each non-root equation is eliminated exactly once, so the + * reduction factors can be stored in O(n) additional persistent storage per + * batch system, rather than PCR's O(n log n) trajectory. + * + * Each batch system is assigned one Kokkos team. The team size is selected + * using Kokkos::AUTO; equations are distributed across team members using + * strided ("grid-stride") loops, so the implementation does not depend on + * the team size being equal to the matrix dimension. + * + * Persistent meaning of the stored arrays after setup(): + * + * - For every index i that is NOT the root: + * main_diagonal(i) = 1 / b_i (inverse pivot at elimination) + * sub_diagonal(i) = a_i / b_i (normalized left coupling, "qLeft") + * q_right_trajectory(i) = c_i / b_i (normalized right coupling, "qRight") + * where a_i, b_i, c_i are the coefficients of equation i at the CR level + * at which i was eliminated. + * + * - For the root index (root_index_): + * main_diagonal(root) = final reduced diagonal (NOT inverted). + * + * The sub-diagonal array is reused for both the "a" (left) and "c" (right) + * coupling of the original problem via the symmetry a_i = c_{i-1}; only one + * off-diagonal array is required, consistent with BatchedTridiagonalSolverThomas + * and BatchedTridiagonalSolverPCR. + * + * For cyclic systems, the same Sherman-Morrison rank-one correction used by + * Thomas/PCR is applied before CR factorization. Because CR overwrites the + * sub-diagonal array in place, the original cyclic corner coefficient is + * copied to cyclic_corner_ once, before any CR level runs, and is never + * re-derived from sub_diagonal_ afterwards. + * + * setup() performs the coefficient factorization once; solve() may be + * called any number of times afterward and never modifies the persistent + * factorization (main_diagonal_, sub_diagonal_, q_right_trajectory_, + * gamma_, cyclic_corner_). + */ + +#include +#include +#include +#include + +#include "../../LinearAlgebra/Vector/vector.h" +#include "../../LinearAlgebra/Vector/vector_operations.h" + +namespace gmgpolar +{ + +template +class BatchedTridiagonalSolverCR : public BatchedTridiagonalSolverBase +{ +public: + BatchedTridiagonalSolverCR(int matrix_dimension, int batch_count, bool is_cyclic = true) + : BatchedTridiagonalSolverBase(matrix_dimension, batch_count, is_cyclic) + , num_steps_(compute_num_steps(matrix_dimension)) + , root_index_(matrix_dimension > 1 ? ((1 << num_steps_) - 1) : 0) + , q_right_trajectory_("BatchedTridiagonalSolverCR::q_right_trajectory", + static_cast(matrix_dimension) * static_cast(batch_count)) + , cyclic_corner_("BatchedTridiagonalSolverCR::cyclic_corner", is_cyclic ? batch_count : 0) + , gamma_("BatchedTridiagonalSolverCR::gamma", is_cyclic ? batch_count : 0) + , is_factorized_(false) + { + assign(q_right_trajectory_, T(0)); + assign(cyclic_corner_, T(0)); + assign(gamma_, T(0)); + } + + /** + * @brief Returns the stored normalized right-coupling factor q_right(i) = c_i / b_i + * for an equation eliminated during setup(). + * + * Only meaningful for indices that are not the root index. + * + * @param batch_idx Batch index. + * @param index Original equation index. + */ + KOKKOS_INLINE_FUNCTION const T& q_right(int batch_idx, int index) const + { + return q_right_trajectory_(static_cast(batch_idx) * this->matrix_dimension_ + index); + } + + /* -------------------------------------------------------- */ + /* Setup: Classical Cyclic Reduction coefficient elimination */ + /* -------------------------------------------------------- */ + // Performs, for cyclic systems, the Sherman-Morrison diagonal adjustment + // (and saves the original corner coefficient), then performs the CR + // level-by-level coefficient reduction down to a single root equation. + + void setup() override + { + int matrix_dimension = this->matrix_dimension_; + int num_steps = num_steps_; + int root_index = root_index_; + bool is_cyclic = this->is_cyclic_; + + Vector main_diagonal = this->main_diagonal_; + Vector sub_diagonal = this->sub_diagonal_; + Vector q_right_trajectory = q_right_trajectory_; + Vector gamma = gamma_; + Vector cyclic_corner = cyclic_corner_; + + (void)root_index; // root requires no explicit action in setup(); it is + // simply whichever index survives every CR level. + + using TeamPolicy = Kokkos::TeamPolicy; + using TeamMember = typename TeamPolicy::member_type; + + TeamPolicy policy(this->batch_count_, Kokkos::AUTO); + + Kokkos::parallel_for( + "SetupCR", policy, KOKKOS_LAMBDA(const TeamMember& team_member) { + const int batch_idx = team_member.league_rank(); + const int offset = batch_idx * matrix_dimension; + const int team_size = team_member.team_size(); + const int rank = team_member.team_rank(); + + // ------------------------------------------------- // + // Sherman-Morrison Adjustment (cyclic systems only) // + // - Save the original cyclic corner coefficient // + // BEFORE it can be overwritten by CR. // + // - Modify the first and last main diagonal element. // + // - Compute and store gamma for later use. // + // ------------------------------------------------- // + if (is_cyclic) { + if (rank == 0) { + const T corner_element = sub_diagonal(offset + matrix_dimension - 1); + cyclic_corner(batch_idx) = corner_element; + gamma(batch_idx) = -main_diagonal(offset + 0); + main_diagonal(offset + 0) -= gamma(batch_idx); + main_diagonal(offset + matrix_dimension - 1) -= + corner_element * corner_element / gamma(batch_idx); + } + team_member.team_barrier(); + } + + // ------------------------------- // + // Classical Cyclic Reduction levels // + // ------------------------------- // + for (int step = 0; step < num_steps; ++step) { + const int d = 1 << step; + + // ---------------------------------------------------- // + // Phase A: eliminate equations i = (d-1) + 2*d*k. // + // Read current (raw) coefficients, store normalized // + // elimination factors in place. // + // ---------------------------------------------------- // + for (int k = rank;; k += team_size) { + const int i = (d - 1) + 2 * d * k; + if (i >= matrix_dimension) { + break; + } + + const bool has_left = (i - d) >= 0; + const bool has_right = (i + d) < matrix_dimension; + + const T b_i = main_diagonal(offset + i); + const T a_i = has_left ? sub_diagonal(offset + i - d) : T(0); + const T c_i = has_right ? sub_diagonal(offset + i) : T(0); + + const T inv_b = T(1) / b_i; + const T qLeft = a_i * inv_b; + const T qRight = c_i * inv_b; + + main_diagonal(offset + i) = inv_b; + sub_diagonal(offset + i) = qLeft; + q_right_trajectory(offset + i) = qRight; + } + + team_member.team_barrier(); + + // ---------------------------------------------------- // + // Phase B: update equations surviving to the next // + // level, s = (2*d-1) + 2*d*k, using ONLY the just- // + // stored normalized factors of the flanking eliminated // + // equations (eL = s-d always exists; eR = s+d may not).// + // ---------------------------------------------------- // + for (int k = rank;; k += team_size) { + const int s = (2 * d - 1) + 2 * d * k; + if (s >= matrix_dimension) { + break; + } + + const int eL = s - d; + const int eR = s + d; + const bool has_eR = eR < matrix_dimension; + + const T qRightEL = q_right_trajectory(offset + eL); + const T b_eL = T(1) / main_diagonal(offset + eL); + + const T c_s_raw = sub_diagonal(offset + s); + const T b_s_raw = main_diagonal(offset + s); + + const T term1 = qRightEL * qRightEL * b_eL; + + T term2 = T(0); + T c_new = T(0); + if (has_eR) { + const T qLeftER = sub_diagonal(offset + eR); // already qLeft(eR) post Phase A + const T qRightER = q_right_trajectory(offset + eR); + term2 = qLeftER * c_s_raw; + c_new = -c_s_raw * qRightER; + } + + main_diagonal(offset + s) = b_s_raw - term1 - term2; + sub_diagonal(offset + s) = c_new; + } + + team_member.team_barrier(); + } + // After the loop, main_diagonal(offset + root_index) holds the + // final reduced (non-inverted) root diagonal. + }); + + Kokkos::fence(); + is_factorized_ = true; + } + + /* ---------------------------------------------------- */ + /* Solve: CR forward RHS reduction + backward substitution */ + /* ---------------------------------------------------- */ + // Reuses the factorization from setup(). Eliminated-equation RHS values + // are preserved (never overwritten) during forward reduction so they + // remain available for backward substitution. + + void solve(Vector rhs, int batch_offset = 0, int batch_stride = 1) override + { + if (!is_factorized_) { + throw std::runtime_error("Error: Matrix must be factorized before solving."); + } + if (batch_stride <= 0) { + throw std::invalid_argument("Error: batch_stride must be positive."); + } + if (batch_offset < 0 || batch_offset > this->batch_count_) { + throw std::invalid_argument("Error: batch_offset out of range."); + } + + int effective_batch_count = (this->batch_count_ - batch_offset + batch_stride - 1) / batch_stride; + if (effective_batch_count < 0) { + effective_batch_count = 0; + } + + int matrix_dimension = this->matrix_dimension_; + int num_steps = num_steps_; + int root_index = root_index_; + bool is_cyclic = this->is_cyclic_; + + Vector main_diagonal = this->main_diagonal_; + Vector sub_diagonal = this->sub_diagonal_; + Vector q_right_trajectory = q_right_trajectory_; + Vector gamma = gamma_; + Vector cyclic_corner = cyclic_corner_; + + using TeamPolicy = Kokkos::TeamPolicy; + using TeamMember = typename TeamPolicy::member_type; + + if (!is_cyclic) { + TeamPolicy policy(effective_batch_count, Kokkos::AUTO); + + Kokkos::parallel_for( + "SolveCRNonCyclic", policy, KOKKOS_LAMBDA(const TeamMember& team_member) { + const int k = team_member.league_rank(); + const int batch_idx = batch_stride * k + batch_offset; + const int offset = batch_idx * matrix_dimension; + const int team_size = team_member.team_size(); + const int rank = team_member.team_rank(); + + // -------------------- // + // Forward RHS Reduction // + // -------------------- // + for (int step = 0; step < num_steps; ++step) { + const int d = 1 << step; + + for (int kk = rank;; kk += team_size) { + const int s = (2 * d - 1) + 2 * d * kk; + if (s >= matrix_dimension) { + break; + } + + const int eL = s - d; + const int eR = s + d; + const bool has_eR = eR < matrix_dimension; + + const T factor_left = q_right_trajectory(offset + eL); + const T factor_right = has_eR ? sub_diagonal(offset + eR) : T(0); + const T rhs_eR = has_eR ? rhs(offset + eR) : T(0); + + rhs(offset + s) -= factor_left * rhs(offset + eL) + factor_right * rhs_eR; + } + + team_member.team_barrier(); + } + + // ----------- // + // Root Solve // + // ----------- // + if (rank == 0) { + rhs(offset + root_index) /= main_diagonal(offset + root_index); + } + team_member.team_barrier(); + + // ---------------------- // + // Backward Substitution // + // ---------------------- // + for (int step = num_steps - 1; step >= 0; --step) { + const int d = 1 << step; + + for (int kk = rank;; kk += team_size) { + const int i = (d - 1) + 2 * d * kk; + if (i >= matrix_dimension) { + break; + } + + const bool has_left = (i - d) >= 0; + const bool has_right = (i + d) < matrix_dimension; + + const T x_left = has_left ? rhs(offset + i - d) : T(0); + const T x_right = has_right ? rhs(offset + i + d) : T(0); + + rhs(offset + i) = rhs(offset + i) * main_diagonal(offset + i) - + sub_diagonal(offset + i) * x_left - + q_right_trajectory(offset + i) * x_right; + } + + team_member.team_barrier(); + } + }); + } + else { + // Cyclic solve: reduce the RHS and the Sherman-Morrison auxiliary + // vector simultaneously, using the same stored CR factors. The + // auxiliary vector uses per-team scratch memory rather than a + // second permanent batch_count * n allocation. + const std::size_t scratch_bytes = static_cast(matrix_dimension) * sizeof(T); + + TeamPolicy policy(effective_batch_count, Kokkos::AUTO); + policy.set_scratch_size(0, Kokkos::PerTeam(static_cast(scratch_bytes))); + + Kokkos::parallel_for( + "SolveCRCyclic", policy, KOKKOS_LAMBDA(const TeamMember& team_member) { + const int k = team_member.league_rank(); + const int batch_idx = batch_stride * k + batch_offset; + const int offset = batch_idx * matrix_dimension; + const int team_size = team_member.team_size(); + const int rank = team_member.team_rank(); + + T* buffer = static_cast(team_member.team_scratch(0).get_shmem(scratch_bytes)); + + // ------------------------------------------------ // + // Initialize the Sherman-Morrison auxiliary vector: // + // u = gamma * e_0 + corner * e_{n-1} // + // ------------------------------------------------ // + for (int i = rank; i < matrix_dimension; i += team_size) { + buffer[i] = T(0); + } + team_member.team_barrier(); + if (rank == 0) { + buffer[0] = gamma(batch_idx); + if (matrix_dimension > 1) { + buffer[matrix_dimension - 1] = cyclic_corner(batch_idx); + } + } + team_member.team_barrier(); + + // -------------------- // + // Forward RHS Reduction // + // -------------------- // + for (int step = 0; step < num_steps; ++step) { + const int d = 1 << step; + + for (int kk = rank;; kk += team_size) { + const int s = (2 * d - 1) + 2 * d * kk; + if (s >= matrix_dimension) { + break; + } + + const int eL = s - d; + const int eR = s + d; + const bool has_eR = eR < matrix_dimension; + + const T factor_left = q_right_trajectory(offset + eL); + const T factor_right = has_eR ? sub_diagonal(offset + eR) : T(0); + const T rhs_eR = has_eR ? rhs(offset + eR) : T(0); + const T buf_eR = has_eR ? buffer[eR] : T(0); + + rhs(offset + s) -= factor_left * rhs(offset + eL) + factor_right * rhs_eR; + buffer[s] -= factor_left * buffer[eL] + factor_right * buf_eR; + } + + team_member.team_barrier(); + } + + // ----------- // + // Root Solve // + // ----------- // + if (rank == 0) { + rhs(offset + root_index) /= main_diagonal(offset + root_index); + buffer[root_index] /= main_diagonal(offset + root_index); + } + team_member.team_barrier(); + + // ---------------------- // + // Backward Substitution // + // ---------------------- // + for (int step = num_steps - 1; step >= 0; --step) { + const int d = 1 << step; + + for (int kk = rank;; kk += team_size) { + const int i = (d - 1) + 2 * d * kk; + if (i >= matrix_dimension) { + break; + } + + const bool has_left = (i - d) >= 0; + const bool has_right = (i + d) < matrix_dimension; + + const T x_left = has_left ? rhs(offset + i - d) : T(0); + const T x_right = has_right ? rhs(offset + i + d) : T(0); + const T v_left = has_left ? buffer[i - d] : T(0); + const T v_right = has_right ? buffer[i + d] : T(0); + + const T invB = main_diagonal(offset + i); + const T qLeft = sub_diagonal(offset + i); + const T qRight = q_right_trajectory(offset + i); + + rhs(offset + i) = rhs(offset + i) * invB - qLeft * x_left - qRight * x_right; + buffer[i] = buffer[i] * invB - qLeft * v_left - qRight * v_right; + } + + team_member.team_barrier(); + } + + // ------------------------------- // + // Sherman-Morrison Reconstruction // + // ------------------------------- // + const T corner = cyclic_corner(batch_idx); + const T g = gamma(batch_idx); + + const T dot_product_x_v = rhs(offset + 0) + corner / g * rhs(offset + matrix_dimension - 1); + const T dot_product_u_v = buffer[0] + corner / g * buffer[matrix_dimension - 1]; + const T factor = dot_product_x_v / (T(1) + dot_product_u_v); + + for (int i = rank; i < matrix_dimension; i += team_size) { + rhs(offset + i) -= factor * buffer[i]; + } + }); + } + Kokkos::fence(); + } + + /* ---------------------------- */ + /* Solve: Diagonal Scaling Only */ + /* ---------------------------- */ + // Valid when the underlying matrix has zero off-diagonal coupling. Does + // NOT perform CR reduction. Because CR's persistent storage inverts the + // diagonal at every eliminated (non-root) index, the true diagonal b_i is + // reconstructed as 1/main_diagonal(i) there, and used directly (raw) at + // the root index. For cyclic systems, the same gamma correction used by + // Thomas/PCR is applied to index 0. + + void solve_diagonal(Vector rhs, int batch_offset = 0, int batch_stride = 1) override + { + if (!is_factorized_) { + throw std::runtime_error("Error: Matrix must be factorized before solving."); + } + if (batch_stride <= 0) { + throw std::invalid_argument("Error: batch_stride must be positive."); + } + if (batch_offset < 0 || batch_offset > this->batch_count_) { + throw std::invalid_argument("Error: batch_offset out of range."); + } + + int effective_batch_count = (this->batch_count_ - batch_offset + batch_stride - 1) / batch_stride; + if (effective_batch_count < 0) { + effective_batch_count = 0; + } + + const int matrix_dimension = this->matrix_dimension_; + const int root_index = root_index_; + const bool is_cyclic = this->is_cyclic_; + Vector main_diagonal = this->main_diagonal_; + Vector gamma = gamma_; + + using MDPolicy = Kokkos::MDRangePolicy>; + MDPolicy policy({0, 0}, {effective_batch_count, matrix_dimension}); + + Kokkos::parallel_for( + "SolveDiagonalCR", policy, KOKKOS_LAMBDA(const int k, const int i) { + const int batch_idx = batch_stride * k + batch_offset; + const int offset = batch_idx * matrix_dimension; + + const T b_i = (i == root_index) ? main_diagonal(offset + i) : T(1) / main_diagonal(offset + i); + + if (is_cyclic && i == 0) { + rhs(offset + 0) /= (b_i + gamma(batch_idx)); + } + else { + rhs(offset + i) /= b_i; + } + }); + + Kokkos::fence(); + } + +private: + static int compute_num_steps(int matrix_dimension) + { + // floor(log2(matrix_dimension)) for matrix_dimension > 1, else 0. + // Uses integer bit logic only (no std::log2), per repository policy. + if (matrix_dimension <= 1) { + return 0; + } + int steps = 0; + while ((1 << (steps + 1)) <= matrix_dimension) { + ++steps; + } + return steps; + } + + int num_steps_; // Number of CR reduction levels: floor(log2(n)), 0 for n<=1. + int root_index_; // Index of the single equation surviving all CR levels. + + Vector q_right_trajectory_; // Persistent normalized right-coupling factors, O(batch_count * n). + + Vector cyclic_corner_; // Original cyclic corner coefficient, saved before CR overwrites sub_diagonal_. + Vector gamma_; // Sherman-Morrison correction factor, one per batch (cyclic systems only). + + bool is_factorized_; +}; + +} // namespace gmgpolar \ No newline at end of file diff --git a/include/LinearAlgebra/Solvers/tridiagonal_solver_crpcr.h b/include/LinearAlgebra/Solvers/tridiagonal_solver_crpcr.h new file mode 100644 index 00000000..65f39f8c --- /dev/null +++ b/include/LinearAlgebra/Solvers/tridiagonal_solver_crpcr.h @@ -0,0 +1,833 @@ +#pragma once + +/** + * @brief Batched tridiagonal solver using hybrid Cyclic Reduction (CR) + + * Parallel Cyclic Reduction (PCR). + * + * This solver is built directly on top of BatchedTridiagonalSolverCR's + * in-place CR factorization/backward-substitution algebra and storage + * model, and on top of BatchedTridiagonalSolverPCR's compact reduction + * recurrence. It is not a redesign of either: it runs the *existing* CR + * levels while the active system is large, and switches to the *existing* + * PCR recurrence -- restricted to a small compact survivor system -- once + * the number of surviving equations is small. + * + * ------------------------------------------------------------------------- + * Algorithm + * ------------------------------------------------------------------------- + * + * original system (size n) + * | + * v + * CR forward reduction, in place, exactly as in BatchedTridiagonalSolverCR + * | (stops after L levels, as soon as the number of CR survivors + * | m satisfies m <= PCR_TARGET_SIZE) + * v + * compact survivor system (size m <= 128) + * | + * v + * compact PCR factorization/solve (team scratch, O(m) per buffer) + * | + * v + * survivor solutions scattered back into their original CR positions + * | + * v + * CR backward substitution, in place, exactly as in BatchedTridiagonalSolverCR + * | + * v + * full solution + */ + +#include +#include +#include + +#include "../../LinearAlgebra/Vector/vector.h" +#include "../../LinearAlgebra/Vector/vector_operations.h" + +namespace gmgpolar +{ + +template +class BatchedTridiagonalSolverCRPCR : public BatchedTridiagonalSolverBase +{ +public: + // Fixed PCR target: CR runs until the survivor count is <= this value. + // No runtime cutoff parameter exists; this constant is isolated here + // so a future version can change it without touching the CR logic. + static constexpr int PCR_TARGET_SIZE = 128; + + BatchedTridiagonalSolverCRPCR(int matrix_dimension, int batch_count, bool is_cyclic = true) + : BatchedTridiagonalSolverBase(matrix_dimension, batch_count, is_cyclic) + , num_cr_levels_(compute_num_cr_levels(matrix_dimension)) + , survivor_base_(num_cr_levels_ > 0 ? ((1 << num_cr_levels_) - 1) : 0) + , survivor_stride_(num_cr_levels_ > 0 ? (1 << num_cr_levels_) : 1) + , m_(compute_survivor_count(matrix_dimension, num_cr_levels_)) + , pcr_num_steps_(compute_pcr_steps(m_)) + , q_right_trajectory_("BatchedTridiagonalSolverCRPCR::q_right_trajectory", + static_cast(matrix_dimension) * static_cast(batch_count)) + , cyclic_corner_("BatchedTridiagonalSolverCRPCR::cyclic_corner", is_cyclic ? batch_count : 0) + , gamma_("BatchedTridiagonalSolverCRPCR::gamma", is_cyclic ? batch_count : 0) + , pcr_k1_trajectory_("BatchedTridiagonalSolverCRPCR::pcr_k1_trajectory", + static_cast(batch_count) * static_cast(m_) * + static_cast(pcr_num_steps_)) + , pcr_k2_trajectory_("BatchedTridiagonalSolverCRPCR::pcr_k2_trajectory", + static_cast(batch_count) * static_cast(m_) * + static_cast(pcr_num_steps_)) + , is_factorized_(false) + { + assign(q_right_trajectory_, T(0)); + assign(cyclic_corner_, T(0)); + assign(gamma_, T(0)); + assign(pcr_k1_trajectory_, T(0)); + assign(pcr_k2_trajectory_, T(0)); + } + + // Number of CR levels actually executed (0 when matrix_dimension <= PCR_TARGET_SIZE). + int crLevels() const + { + return num_cr_levels_; + } + + // Number of CR survivors / compact PCR dimension (<= PCR_TARGET_SIZE). + int survivorCount() const + { + return m_; + } + + // Number of compact PCR reduction steps needed to reduce m equations, 0 if m<=1. + int pcrSteps() const + { + return pcr_num_steps_; + } + + /* -------------------------------------------------------------- */ + /* Setup: CR forward factorization (existing algebra) down to the */ + /* survivor boundary, then compact PCR factorization of the m<=128 */ + /* survivor system, writing the final diagonal back into the */ + /* survivor slots of main_diagonal_. */ + /* -------------------------------------------------------------- */ + void setup() override + { + const int n = this->matrix_dimension_; + const int L = num_cr_levels_; + const int survivor_base = survivor_base_; + const int survivor_stride = survivor_stride_; + const int m = m_; + const int pcr_steps = pcr_num_steps_; + const bool is_cyclic = this->is_cyclic_; + + Vector main_diagonal = this->main_diagonal_; + Vector sub_diagonal = this->sub_diagonal_; + Vector q_right_trajectory = q_right_trajectory_; + Vector gamma = gamma_; + Vector cyclic_corner = cyclic_corner_; + Vector pcr_k1 = pcr_k1_trajectory_; + Vector pcr_k2 = pcr_k2_trajectory_; + + using TeamPolicy = Kokkos::TeamPolicy; + using TeamMember = typename TeamPolicy::member_type; + + // n == 1: the cyclic-corner storage slot IS index 0 + // (matrix_dimension - 1 == 0), so the generic Sherman-Morrison + // derivation is not meaningful. BatchedTridiagonalSolverPCR + // explicitly special-cases matrix_dimension == 1 the same way: + // it is treated as a plain 1x1 diagonal system regardless of + // is_cyclic. CRPCR matches that established convention exactly + // (verified against a dense reference in the standalone harness). + if (n == 1) { + is_factorized_ = true; + return; + } + + const std::size_t scratch_bytes = 4ull * static_cast(m) * sizeof(T); + + TeamPolicy policy(this->batch_count_, Kokkos::AUTO); + policy.set_scratch_size(0, Kokkos::PerTeam(scratch_bytes)); + + Kokkos::parallel_for( + "SetupCRPCR", policy, KOKKOS_LAMBDA(const TeamMember& team_member) { + const int batch_idx = team_member.league_rank(); + const int offset = batch_idx * n; + const int team_size = team_member.team_size(); + const int rank = team_member.team_rank(); + + // -------------------------------------------------- // + // Sherman-Morrison Adjustment (cyclic systems only), // + // identical to BatchedTridiagonalSolverCR::setup(). // + // Saved BEFORE CR can overwrite sub_diagonal_. // + // -------------------------------------------------- // + if (is_cyclic) { + if (rank == 0) { + const T corner_element = sub_diagonal(offset + n - 1); + cyclic_corner(batch_idx) = corner_element; + gamma(batch_idx) = -main_diagonal(offset + 0); + main_diagonal(offset + 0) -= gamma(batch_idx); + main_diagonal(offset + n - 1) -= corner_element * corner_element / gamma(batch_idx); + } + team_member.team_barrier(); + } + + // --------------------------------------------------- // + // CR forward levels 0 .. L-1: IDENTICAL arithmetic to // + // BatchedTridiagonalSolverCR::setup(), just bounded // + // to L levels instead of running to a single root. // + // --------------------------------------------------- // + for (int step = 0; step < L; ++step) { + const int d = 1 << step; + + // Phase A: eliminate i = (d-1) + 2*d*k. + for (int k = rank;; k += team_size) { + const int i = (d - 1) + 2 * d * k; + if (i >= n) { + break; + } + + const bool has_left = (i - d) >= 0; + const bool has_right = (i + d) < n; + + const T b_i = main_diagonal(offset + i); + const T a_i = has_left ? sub_diagonal(offset + i - d) : T(0); + const T c_i = has_right ? sub_diagonal(offset + i) : T(0); + + const T inv_b = T(1) / b_i; + const T qLeft = a_i / b_i; + const T qRight = c_i / b_i; + + main_diagonal(offset + i) = inv_b; + sub_diagonal(offset + i) = qLeft; + q_right_trajectory(offset + i) = qRight; + } + + team_member.team_barrier(); + + // Phase B: survivors s = (2*d-1) + 2*d*k. + for (int k = rank;; k += team_size) { + const int s = (2 * d - 1) + 2 * d * k; + if (s >= n) { + break; + } + + const int eL = s - d; + const int eR = s + d; + const bool has_eR = eR < n; + + const T qRightEL = q_right_trajectory(offset + eL); + + const T c_s_raw = sub_diagonal(offset + s); + const T b_s_raw = main_diagonal(offset + s); + + const T term1 = qRightEL * qRightEL / main_diagonal(offset + eL); + + T term2 = T(0); + T c_new = T(0); + if (has_eR) { + const T qLeftER = sub_diagonal(offset + eR); + const T qRightER = q_right_trajectory(offset + eR); + term2 = qLeftER * c_s_raw; + c_new = -c_s_raw * qRightER; + } + + main_diagonal(offset + s) = b_s_raw - term1 - term2; + sub_diagonal(offset + s) = c_new; + } + + team_member.team_barrier(); + } + + // ------------------------------------------------- // + // Compact gather: survivors r_j = survivor_base + // + // j*survivor_stride hold "current reduced diagonal" // + // (main_diagonal) and "current reduced right // + // coefficient" (sub_diagonal), per the persistent // + // storage contract. // + // ------------------------------------------------- // + T* scratch = static_cast(team_member.team_scratch(0).get_shmem(scratch_bytes)); + // Layout: [e0 | b0 | e1 | b1], each length m. + T* e[2] = {scratch, scratch + 2 * m}; + T* b[2] = {scratch + m, scratch + 3 * m}; + + for (int j = rank; j < m; j += team_size) { + const int orig = survivor_base + j * survivor_stride; + b[0][j] = main_diagonal(offset + orig); + e[0][j] = sub_diagonal(offset + orig); + } + team_member.team_barrier(); + + // When L==0 (n<=PCR_TARGET_SIZE: no CR ran), the survivor + // set is the raw original system and the last equation's + // sub_diagonal_ slot is either unused (non-cyclic) or holds + // the cyclic corner (cyclic) -- never a real right-neighbor + // coupling. Zero it explicitly, exactly as + // BatchedTridiagonalSolverPCR::setup() already does + // (e[cur][i] = (i==n-1) ? 0 : sub_diagonal(...)). When + // L>=1 this is already guaranteed by CR's own Phase B + // (c_new==0 whenever the survivor has no right CR + // neighbor, which is exactly true for the last survivor), + // so no separate action is required in that case. + if (L == 0 && rank == 0 && m > 0) { + e[0][m - 1] = T(0); + } + team_member.team_barrier(); + + // ----------------------------------------------------- // + // Compact PCR factorization: IDENTICAL recurrence to // + // BatchedTridiagonalSolverPCR::setup(), re-scoped from // + // size n to size m. Symmetry is used: only e[] (the // + // right coefficient) is stored per level; the left // + // coefficient a(j) = c(j-delta) = e[cur][j-delta] is // + // reconstructed, never stored separately. // + // Two multiplier trajectories (k1, k2) are stored // + // because k1(j) != k2(j-delta) in general (Section 6): // + // this mirrors the already-proven-correct convention of // + // the reference PCR solver rather than inventing a new, // + // unproven single-array compression. // + // ----------------------------------------------------- // + int cur = 0; + for (int step = 0; step < pcr_steps; ++step) { + const int delta = 1 << step; + + for (int j = rank; j < m; j += team_size) { + const bool has_left = j >= delta; + const bool has_right = j + delta < m; + const int left = j - delta; + const int right = j + delta; + + // The compact system is symmetric. Therefore the left + // coefficient a(j) is c(j-1) and, at this PCR level, + // a(j) = c(j-delta). Boundary terms must be zero rather + // than clamped to an existing equation. + const T a_j = has_left ? e[cur][left] : T(0); + const T c_j = has_right ? e[cur][j] : T(0); + const T c_right = has_right ? e[cur][right] : T(0); + + const T k1_val = has_left ? a_j / b[cur][left] : T(0); + const T k2_val = has_right ? c_j / b[cur][right] : T(0); + + const std::size_t traj = static_cast(batch_idx) * + static_cast(pcr_steps) * static_cast(m) + + static_cast(step) * static_cast(m) + + static_cast(j); + + pcr_k1(traj) = k1_val; + pcr_k2(traj) = k2_val; + + const int nxt = 1 - cur; + e[nxt][j] = has_right ? -c_right * k2_val : T(0); + b[nxt][j] = b[cur][j] - a_j * k1_val - c_j * k2_val; + } + + team_member.team_barrier(); + cur = 1 - cur; + } + + // ---------------------------------------------------- // + // Write the final PCR diagonal back into the survivor // + // positions of main_diagonal_ (Section 10). No // + // persistent compact pcr_main_diagonal_ array is // + // created. The final PCR off-diagonal is not needed // + // by solve() and is not persisted (Section 11). // + // Eliminated (non-survivor) main_diagonal_ entries are // + // untouched by this loop. // + // ---------------------------------------------------- // + for (int j = rank; j < m; j += team_size) { + const int orig = survivor_base + j * survivor_stride; + main_diagonal(offset + orig) = b[cur][j]; + } + }); + + Kokkos::fence(); + is_factorized_ = true; + } + + /* --------------------------------------------------------------- */ + /* Solve: CR forward RHS reduction (existing algebra) -> compact */ + /* survivor gather -> compact PCR solve (replay of the stored */ + /* factorization) -> scatter -> CR backward substitution (existing */ + /* algebra, unchanged). */ + /* --------------------------------------------------------------- */ + void solve(Vector rhs, int batch_offset = 0, int batch_stride = 1) override + { + if (!is_factorized_) { + throw std::runtime_error("Error: Matrix must be factorized before solving."); + } + if (batch_stride <= 0) { + throw std::invalid_argument("Error: batch_stride must be positive."); + } + if (batch_offset < 0 || batch_offset > this->batch_count_) { + throw std::invalid_argument("Error: batch_offset out of range."); + } + + const int effective_batch_count = compute_effective_batch_count(batch_offset, batch_stride); + if (effective_batch_count == 0) { + return; + } + + const int n = this->matrix_dimension_; + const int L = num_cr_levels_; + const int survivor_base = survivor_base_; + const int survivor_stride = survivor_stride_; + const int m = m_; + const int pcr_steps = pcr_num_steps_; + const bool is_cyclic = this->is_cyclic_; + + Vector main_diagonal = this->main_diagonal_; + Vector sub_diagonal = this->sub_diagonal_; + Vector q_right_trajectory = q_right_trajectory_; + Vector gamma = gamma_; + Vector cyclic_corner = cyclic_corner_; + Vector pcr_k1 = pcr_k1_trajectory_; + Vector pcr_k2 = pcr_k2_trajectory_; + + using TeamPolicy = Kokkos::TeamPolicy; + using TeamMember = typename TeamPolicy::member_type; + + if (n == 1) { + // Matches BatchedTridiagonalSolverPCR's n==1 special case: + // plain diagonal solve, regardless of is_cyclic. + Kokkos::parallel_for( + "SolveCRPCRTrivial", Kokkos::RangePolicy(0, effective_batch_count), + KOKKOS_LAMBDA(const int k) { + const int batch_idx = batch_stride * k + batch_offset; + rhs(batch_idx) /= main_diagonal(batch_idx); + }); + Kokkos::fence(); + return; + } + + if (!is_cyclic) { + const std::size_t scratch_bytes = 2ull * static_cast(m) * sizeof(T); + + TeamPolicy policy(effective_batch_count, Kokkos::AUTO); + policy.set_scratch_size(0, Kokkos::PerTeam(scratch_bytes)); + + Kokkos::parallel_for( + "SolveCRPCRNonCyclic", policy, KOKKOS_LAMBDA(const TeamMember& team_member) { + const int k = team_member.league_rank(); + const int batch_idx = batch_stride * k + batch_offset; + const int offset = batch_idx * n; + const int team_size = team_member.team_size(); + const int rank = team_member.team_rank(); + + // -------------------- // + // CR Forward RHS Reduction (levels 0..L-1), identical // + // algebra to BatchedTridiagonalSolverCR::solve(). // + // -------------------- // + for (int step = 0; step < L; ++step) { + const int d = 1 << step; + + for (int kk = rank;; kk += team_size) { + const int s = (2 * d - 1) + 2 * d * kk; + if (s >= n) { + break; + } + + const int eL = s - d; + const int eR = s + d; + const bool has_eR = eR < n; + + const T factor_left = q_right_trajectory(offset + eL); + const T factor_right = has_eR ? sub_diagonal(offset + eR) : T(0); + const T rhs_eR = has_eR ? rhs(offset + eR) : T(0); + + rhs(offset + s) -= factor_left * rhs(offset + eL) + factor_right * rhs_eR; + } + + team_member.team_barrier(); + } + + // ------------- // + // Compact survivor gather // + // ------------- // + T* scratch = static_cast(team_member.team_scratch(0).get_shmem(scratch_bytes)); + T* d[2] = {scratch, scratch + m}; + + int cur = 0; + for (int j = rank; j < m; j += team_size) { + const int orig = survivor_base + j * survivor_stride; + d[cur][j] = rhs(offset + orig); + } + team_member.team_barrier(); + + // ----------------------------------------------------- // + // Compact PCR solve (trajectory replay), identical to // + // BatchedTridiagonalSolverPCR::solve(), re-scoped to m. // + // ----------------------------------------------------- // + for (int step = 0; step < pcr_steps; ++step) { + const int delta = 1 << step; + + for (int j = rank; j < m; j += team_size) { + const bool has_left = j >= delta; + const bool has_right = j + delta < m; + const int left = j - delta; + const int right = j + delta; + + const std::size_t traj = static_cast(batch_idx) * + static_cast(pcr_steps) * + static_cast(m) + + static_cast(step) * static_cast(m) + + static_cast(j); + + const T k1_val = pcr_k1(traj); + const T k2_val = pcr_k2(traj); + + const int nxt = 1 - cur; + + d[nxt][j] = d[cur][j] - (has_left ? k1_val * d[cur][left] : T(0)) - + (has_right ? k2_val * d[cur][right] : T(0)); + } + + team_member.team_barrier(); + cur = 1 - cur; + } + + // PCR is a complete solve: divide by the final diagonal + // and scatter directly into the original survivor + // positions. No PCR backward-substitution phase. + for (int j = rank; j < m; j += team_size) { + const int orig = survivor_base + j * survivor_stride; + rhs(offset + orig) = d[cur][j] / main_diagonal(offset + orig); + } + team_member.team_barrier(); + + // --------------------------------------------------- // + // CR Backward Substitution (levels L-1..0), identical // + // algebra to BatchedTridiagonalSolverCR::solve(). // + // --------------------------------------------------- // + for (int step = L - 1; step >= 0; --step) { + const int d = 1 << step; + + for (int kk = rank;; kk += team_size) { + const int i = (d - 1) + 2 * d * kk; + if (i >= n) { + break; + } + + const bool has_left = (i - d) >= 0; + const bool has_right = (i + d) < n; + + const T x_left = has_left ? rhs(offset + i - d) : T(0); + const T x_right = has_right ? rhs(offset + i + d) : T(0); + + rhs(offset + i) = rhs(offset + i) * main_diagonal(offset + i) - + sub_diagonal(offset + i) * x_left - + q_right_trajectory(offset + i) * x_right; + } + + team_member.team_barrier(); + } + }); + } + else { + // Cyclic solve: the actual RHS and the Sherman-Morrison + // auxiliary vector are carried through the identical + // CR-forward -> compact-PCR -> CR-backward pipeline together, + // reusing the single stored factorization, then combined via + // the existing Sherman-Morrison reconstruction. The auxiliary + // vector uses O(n) team scratch (the same pattern already used + // by BatchedTridiagonalSolverCR::solve()'s cyclic branch) + // rather than a permanent batch_count*n array. + const std::size_t aux_bytes = static_cast(n) * sizeof(T); + const std::size_t compact_bytes = 4ull * static_cast(m) * sizeof(T); + const std::size_t scratch_bytes = aux_bytes + compact_bytes; + + TeamPolicy policy(effective_batch_count, Kokkos::AUTO); + policy.set_scratch_size(0, Kokkos::PerTeam(scratch_bytes)); + + Kokkos::parallel_for( + "SolveCRPCRCyclic", policy, KOKKOS_LAMBDA(const TeamMember& team_member) { + const int k = team_member.league_rank(); + const int batch_idx = batch_stride * k + batch_offset; + const int offset = batch_idx * n; + const int team_size = team_member.team_size(); + const int rank = team_member.team_rank(); + + T* raw = static_cast(team_member.team_scratch(0).get_shmem(scratch_bytes)); + T* buffer = raw; // O(n) Sherman-Morrison auxiliary vector. + T* compact = raw + n; // O(4m) compact PCR scratch. + T* d_rhs[2] = {compact, compact + 2 * m}; + T* d_buf[2] = {compact + m, compact + 3 * m}; + + // ------------------------------------------------- // + // Initialize the Sherman-Morrison auxiliary vector: // + // u = gamma * e_0 + corner * e_{n-1}, identical to // + // BatchedTridiagonalSolverCR::solve(). // + // ------------------------------------------------- // + for (int i = rank; i < n; i += team_size) { + buffer[i] = T(0); + } + team_member.team_barrier(); + if (rank == 0) { + buffer[0] = gamma(batch_idx); + if (n > 1) { + buffer[n - 1] = cyclic_corner(batch_idx); + } + } + team_member.team_barrier(); + + // ------------------------------------------------- // + // CR Forward RHS Reduction, actual RHS + auxiliary, // + // identical algebra to BatchedTridiagonalSolverCR. // + // ------------------------------------------------- // + for (int step = 0; step < L; ++step) { + const int d = 1 << step; + + for (int kk = rank;; kk += team_size) { + const int s = (2 * d - 1) + 2 * d * kk; + if (s >= n) { + break; + } + + const int eL = s - d; + const int eR = s + d; + const bool has_eR = eR < n; + + const T factor_left = q_right_trajectory(offset + eL); + const T factor_right = has_eR ? sub_diagonal(offset + eR) : T(0); + const T rhs_eR = has_eR ? rhs(offset + eR) : T(0); + const T buf_eR = has_eR ? buffer[eR] : T(0); + + rhs(offset + s) -= factor_left * rhs(offset + eL) + factor_right * rhs_eR; + buffer[s] -= factor_left * buffer[eL] + factor_right * buf_eR; + } + + team_member.team_barrier(); + } + + // --------------------------------- // + // Compact gather (both quantities). // + // --------------------------------- // + int cur = 0; + for (int j = rank; j < m; j += team_size) { + const int orig = survivor_base + j * survivor_stride; + d_rhs[cur][j] = rhs(offset + orig); + d_buf[cur][j] = buffer[orig]; + } + team_member.team_barrier(); + + // ----------------------------------------------- // + // Compact PCR solve (trajectory replay) for both. // + // ----------------------------------------------- // + for (int step = 0; step < pcr_steps; ++step) { + const int delta = 1 << step; + + for (int j = rank; j < m; j += team_size) { + const bool has_left = j >= delta; + const bool has_right = j + delta < m; + const int left = j - delta; + const int right = j + delta; + + const std::size_t traj = static_cast(batch_idx) * + static_cast(pcr_steps) * + static_cast(m) + + static_cast(step) * static_cast(m) + + static_cast(j); + + const T k1_val = pcr_k1(traj); + const T k2_val = pcr_k2(traj); + + const int nxt = 1 - cur; + d_rhs[nxt][j] = d_rhs[cur][j] - (has_left ? k1_val * d_rhs[cur][left] : T(0)) - + (has_right ? k2_val * d_rhs[cur][right] : T(0)); + d_buf[nxt][j] = d_buf[cur][j] - (has_left ? k1_val * d_buf[cur][left] : T(0)) - + (has_right ? k2_val * d_buf[cur][right] : T(0)); + } + + team_member.team_barrier(); + cur = 1 - cur; + } + + // Scatter both survivor solutions. + for (int j = rank; j < m; j += team_size) { + const int orig = survivor_base + j * survivor_stride; + rhs(offset + orig) = d_rhs[cur][j] / main_diagonal(offset + orig); + buffer[orig] = d_buf[cur][j] / main_diagonal(offset + orig); + } + team_member.team_barrier(); + + // ---------------------------------------------------- // + // CR Backward Substitution for both, identical algebra // + // to BatchedTridiagonalSolverCR::solve(). // + // ---------------------------------------------------- // + for (int step = L - 1; step >= 0; --step) { + const int d = 1 << step; + + for (int kk = rank;; kk += team_size) { + const int i = (d - 1) + 2 * d * kk; + if (i >= n) { + break; + } + + const bool has_left = (i - d) >= 0; + const bool has_right = (i + d) < n; + + const T x_left = has_left ? rhs(offset + i - d) : T(0); + const T x_right = has_right ? rhs(offset + i + d) : T(0); + const T v_left = has_left ? buffer[i - d] : T(0); + const T v_right = has_right ? buffer[i + d] : T(0); + + const T invB = main_diagonal(offset + i); + const T qLeft = sub_diagonal(offset + i); + const T qRight = q_right_trajectory(offset + i); + + rhs(offset + i) = rhs(offset + i) * invB - qLeft * x_left - qRight * x_right; + buffer[i] = buffer[i] * invB - qLeft * v_left - qRight * v_right; + } + + team_member.team_barrier(); + } + + // -------------------------------------------------- // + // Sherman-Morrison Reconstruction, identical formula // + // to BatchedTridiagonalSolverCR::solve(). // + // -------------------------------------------------- // + const T corner = cyclic_corner(batch_idx); + const T g = gamma(batch_idx); + + const T dot_product_x_v = rhs(offset + 0) + corner / g * rhs(offset + n - 1); + const T dot_product_u_v = buffer[0] + corner / g * buffer[n - 1]; + const T factor = dot_product_x_v / (T(1) + dot_product_u_v); + + for (int i = rank; i < n; i += team_size) { + rhs(offset + i) -= factor * buffer[i]; + } + }); + } + Kokkos::fence(); + } + + /* ---------------------------- */ + /* Solve: Diagonal Scaling Only */ + /* ---------------------------- */ + // Does NOT run CR or PCR. Preserves BatchedTridiagonalSolverCR's + // semantics, generalized from "the single CR root" to "the m CR + // survivor positions": eliminated (non-survivor) positions store an + // inverse pivot (1/b_i); survivor positions store the raw final + // diagonal directly (no PCR backward phase exists to undo, matching + // BatchedTridiagonalSolverPCR's own convention for its stored root). + void solve_diagonal(Vector rhs, int batch_offset = 0, int batch_stride = 1) override + { + if (!is_factorized_) { + throw std::runtime_error("Error: Matrix must be factorized before solving."); + } + if (batch_stride <= 0) { + throw std::invalid_argument("Error: batch_stride must be positive."); + } + if (batch_offset < 0 || batch_offset > this->batch_count_) { + throw std::invalid_argument("Error: batch_offset out of range."); + } + + const int effective_batch_count = compute_effective_batch_count(batch_offset, batch_stride); + if (effective_batch_count == 0) { + return; + } + + const int n = this->matrix_dimension_; + const int L = num_cr_levels_; + const int survivor_base = survivor_base_; + const int survivor_stride = survivor_stride_; + const int m = m_; + const bool is_cyclic = this->is_cyclic_; + Vector main_diagonal = this->main_diagonal_; + Vector gamma = gamma_; + + using MDPolicy = Kokkos::MDRangePolicy>; + MDPolicy policy({0, 0}, {effective_batch_count, n}); + + Kokkos::parallel_for( + "SolveDiagonalCRPCR", policy, KOKKOS_LAMBDA(const int k, const int i) { + const int batch_idx = batch_stride * k + batch_offset; + const int offset = batch_idx * n; + + bool is_survivor; + if (L == 0) { + is_survivor = true; + } + else { + is_survivor = (i >= survivor_base) && (((i - survivor_base) % survivor_stride) == 0) && + (((i - survivor_base) / survivor_stride) < m); + } + + const T b_i = is_survivor ? main_diagonal(offset + i) : T(1) / main_diagonal(offset + i); + + if (is_cyclic && i == 0) { + rhs(offset + 0) /= (b_i + gamma(batch_idx)); + } + else { + rhs(offset + i) /= b_i; + } + }); + + Kokkos::fence(); + } + +private: + static int compute_pcr_steps(int survivor_count) + { + int steps = 0; + int stride = 1; + while (stride < survivor_count) { + ++steps; + // survivor_count is bounded by PCR_TARGET_SIZE, so this cannot + // overflow for the configured target. + stride <<= 1; + } + return steps; + } + + int compute_effective_batch_count(int batch_offset, int batch_stride) const + { + if (batch_offset == this->batch_count_) { + return 0; + } + return (this->batch_count_ - batch_offset + batch_stride - 1) / batch_stride; + } + + static int compute_num_cr_levels(int matrix_dimension) + { + if (matrix_dimension <= PCR_TARGET_SIZE) { + return 0; + } + int L = 1; + while (true) { + const long long base = (1LL << L) - 1; + const long long stride = 1LL << L; + if (base >= matrix_dimension) { + return L; // Guard; not expected to trigger for valid n. + } + const long long m = (matrix_dimension - 1 - base) / stride + 1; + if (m <= PCR_TARGET_SIZE) { + return L; + } + ++L; + } + } + + static int compute_survivor_count(int matrix_dimension, int L) + { + if (L == 0) { + return matrix_dimension; + } + const long long base = (1LL << L) - 1; + const long long stride = 1LL << L; + return static_cast((matrix_dimension - 1 - base) / stride + 1); + } + + int num_cr_levels_; // L: number of CR levels executed (0 if n <= PCR_TARGET_SIZE). + int survivor_base_; // 2^L - 1 (0 if L==0). + int survivor_stride_; // 2^L (1 if L==0). + int m_; // CR survivor count / compact PCR dimension, m <= PCR_TARGET_SIZE. + int pcr_num_steps_; // Number of compact PCR reduction steps, 0 if m<=1. + + Vector q_right_trajectory_; // Persistent normalized right-coupling factors for CR-eliminated indices. + + Vector cyclic_corner_; // Original cyclic corner coefficient, saved before CR overwrites sub_diagonal_. + Vector gamma_; // Sherman-Morrison correction factor, one per batch (cyclic systems only). + + Vector pcr_k1_trajectory_; // Compact PCR left-elimination multipliers, O(batch*m*log m), m<=128. + Vector pcr_k2_trajectory_; // Compact PCR right-elimination multipliers, O(batch*m*log m), m<=128. + + bool is_factorized_; +}; + +} // namespace gmgpolar \ No newline at end of file diff --git a/include/LinearAlgebra/Solvers/tridiagonal_solver_pcr.h b/include/LinearAlgebra/Solvers/tridiagonal_solver_pcr.h new file mode 100644 index 00000000..60d3670c --- /dev/null +++ b/include/LinearAlgebra/Solvers/tridiagonal_solver_pcr.h @@ -0,0 +1,477 @@ +#pragma once + +/** + * @brief Batched tridiagonal solver based on Parallel Cyclic Reduction (PCR). + * + * The solver operates on a batch of independent tridiagonal systems and uses + * Kokkos teams to distribute the equations of each system across team members. + * + * PCR performs O(n log n) arithmetic work with O(log n) parallel depth. This + * allows the equations within a system to be processed concurrently, which is + * beneficial when the number of systems is small relative to the system size. + * + * Each system is assigned one Kokkos team. The team size is selected using + * Kokkos::AUTO to allow the execution backend to choose an appropriate value. + * The equations are distributed across team members using strided loops, so + * the implementation does not depend on the team size being equal to the + * matrix dimension. + * + * @tparam T Scalar type used for matrix coefficients and right-hand sides. + * + * @note setup() performs the coefficient reduction and stores the reduction + * coefficients. Subsequent calls to solve() reuse this factorization. + * + * For cyclic systems, setup() also prepares the diagonal correction required + * by the Sherman–Morrison reconstruction. The coefficient reduction is shared + * by the right-hand side and auxiliary solve performed by solve(). + */ + +#include +#include +#include +#include +#include + +#include "../../LinearAlgebra/Vector/vector.h" +#include "../../LinearAlgebra/Vector/vector_operations.h" + +namespace gmgpolar +{ + +/** + * @brief Computes the clamped left and right neighbor indices used by PCR. + * + * The boundary convention is that the left coefficient of the first equation + * and the right coefficient of the last equation are zero. Clamping the + * neighbor indices to the valid range allows the same PCR update to be used + * for boundary and interior equations. + * + * @param i Current equation index. + * @param delta Distance to the neighboring equation for the current PCR step. + * @param n Number of equations. + * @param[out] iLeft Left neighbor index. + * @param[out] iRight Right neighbor index. + */ +KOKKOS_INLINE_FUNCTION +void pcr_neighbors(int i, int delta, int n, int& iLeft, int& iRight) +{ + iLeft = i - delta; + if (iLeft < 0) { + iLeft = 0; + } + + iRight = i + delta; + if (iRight >= n) { + iRight = n - 1; + } +} + +template +class BatchedTridiagonalSolverPCR : public BatchedTridiagonalSolverBase +{ +public: + BatchedTridiagonalSolverPCR(int matrix_dimension, int batch_count, bool is_cyclic = true) + : BatchedTridiagonalSolverBase(matrix_dimension, batch_count, is_cyclic) + , gamma_("BatchedTridiagonalSolverPCR::gamma", is_cyclic ? batch_count : 0) + , is_factorized_(false) + , num_steps_( + matrix_dimension > 1 ? static_cast(std::ceil(std::log2(static_cast(matrix_dimension)))) : 0) + , k1_trajectory_("BatchedTridiagonalSolverPCR::k1_trajectory", static_cast(batch_count) * + static_cast(num_steps_) * + static_cast(matrix_dimension)) + , k2_trajectory_("BatchedTridiagonalSolverPCR::k2_trajectory", static_cast(batch_count) * + static_cast(num_steps_) * + static_cast(matrix_dimension)) + { + assign(gamma_, T(0)); + assign(k1_trajectory_, T(0)); + assign(k2_trajectory_, T(0)); + } + + /** + * @brief Returns a stored PCR left reduction coefficient. + * + * @param batch_idx Batch index. + * @param step PCR reduction step. + * @param index Equation index. + */ + KOKKOS_INLINE_FUNCTION const T& k1(const int batch_idx, const int step, const int index) const + { + return k1_trajectory_(static_cast(batch_idx) * num_steps_ * this->matrix_dimension_ + + static_cast(step) * this->matrix_dimension_ + index); + } + + /** + * @brief Returns a stored PCR right reduction coefficient. + * + * @param batch_idx Batch index. + * @param step PCR reduction step. + * @param index Equation index. + */ + KOKKOS_INLINE_FUNCTION const T& k2(const int batch_idx, const int step, const int index) const + { + return k2_trajectory_(static_cast(batch_idx) * num_steps_ * this->matrix_dimension_ + + static_cast(step) * this->matrix_dimension_ + index); + } + + /** + * @brief Performs PCR coefficient reduction and prepares the solver. + * + * The coefficient reduction is performed once and the resulting reduction + * coefficients are stored for reuse by solve(). The final reduced diagonal + * replaces the original main diagonal. + * + * For cyclic systems, the diagonal is modified as part of the + * Sherman–Morrison formulation and the corresponding correction factor is + * stored in gamma_. + * + * The sub-diagonal remains unchanged by this operation. + */ + void setup() override + { + int matrix_dimension = this->matrix_dimension_; + int num_steps = num_steps_; + bool is_cyclic = this->is_cyclic_; + + Vector main_diagonal = this->main_diagonal_; + Vector sub_diagonal = this->sub_diagonal_; + Vector gamma = gamma_; + Vector k1_trajectory = k1_trajectory_; + Vector k2_trajectory = k2_trajectory_; + + if (matrix_dimension == 1) { + is_factorized_ = true; + return; + } + + using TeamPolicy = Kokkos::TeamPolicy; + using TeamMember = typename TeamPolicy::member_type; + + // For a symmetric tridiagonal system, the left coefficient of an + // equation is the right coefficient of the corresponding left + // neighbor. The left coefficients can therefore be reconstructed from + // e[] rather than stored separately, reducing team scratch storage. + const std::size_t scratch_bytes = 2ull * 2ull * static_cast(matrix_dimension) * sizeof(T); + + TeamPolicy policy(this->batch_count_, Kokkos::AUTO); + policy.set_scratch_size(0, Kokkos::PerTeam(static_cast(scratch_bytes))); + + Kokkos::parallel_for( + "SetupPCR", policy, KOKKOS_LAMBDA(const TeamMember& team_member) { + const int batch_idx = team_member.league_rank(); + const int offset = batch_idx * matrix_dimension; + const int team_size = team_member.team_size(); + const int rank = team_member.team_rank(); + + T* scratch = static_cast(team_member.team_scratch(0).get_shmem(scratch_bytes)); + // Layout: [e0 | b0 | e1 | b1]. + T* e[2] = {scratch, scratch + 2 * matrix_dimension}; + T* b[2] = {scratch + matrix_dimension, scratch + 3 * matrix_dimension}; + + int cur = 0; + + // e[] stores the right coefficients of the current reduced + // system. The left coefficients are reconstructed from e[]. + for (int i = rank; i < matrix_dimension; i += team_size) { + e[cur][i] = (i == matrix_dimension - 1) ? T(0) : sub_diagonal(offset + i); + b[cur][i] = main_diagonal(offset + i); + } + + team_member.team_barrier(); + + if (is_cyclic) { + if (rank == 0) { + const T cyclic_corner_element = sub_diagonal(offset + matrix_dimension - 1); + gamma(batch_idx) = -main_diagonal(offset + 0); + b[cur][0] -= gamma(batch_idx); + b[cur][matrix_dimension - 1] -= + cyclic_corner_element * cyclic_corner_element / gamma(batch_idx); + } + team_member.team_barrier(); + } + + for (int step = 0; step < num_steps; step++) { + const int delta = 1 << step; + + for (int i = rank; i < matrix_dimension; i += team_size) { + int iLeft, iRight; + pcr_neighbors(i, delta, matrix_dimension, iLeft, iRight); + + const T a_i = (i >= delta) ? e[cur][i - delta] : T(0); + const T a_iRight = (iRight >= delta) ? e[cur][iRight - delta] : T(0); + const T c_i = e[cur][i]; + const T c_iLeft = e[cur][iLeft]; + + const T k1_val = a_i / b[cur][iLeft]; + const T k2_val = c_i / b[cur][iRight]; + + k1_trajectory(static_cast(batch_idx) * num_steps * matrix_dimension + + static_cast(step) * matrix_dimension + i) = k1_val; + k2_trajectory(static_cast(batch_idx) * num_steps * matrix_dimension + + static_cast(step) * matrix_dimension + i) = k2_val; + + const int nxt = 1 - cur; + e[nxt][i] = -e[cur][iRight] * k2_val; + b[nxt][i] = b[cur][i] - c_iLeft * k1_val - a_iRight * k2_val; + } + + team_member.team_barrier(); + cur = 1 - cur; + } + + for (int i = rank; i < matrix_dimension; i += team_size) { + main_diagonal(offset + i) = b[cur][i]; + } + }); + + Kokkos::fence(); + is_factorized_ = true; + } + + /** + * @brief Solves the factored tridiagonal systems for the supplied RHS. + * + * The coefficient reduction performed by setup() is reused. The stored PCR + * trajectory is applied to the right-hand side, followed by division by the + * reduced diagonal. + * + * For cyclic systems, the right-hand side and the auxiliary vector used by + * the Sherman–Morrison reconstruction are reduced using the same stored + * trajectory within a single kernel launch. + * + * @param rhs Right-hand sides to solve. The vector is overwritten with the + * corresponding solutions. + * @param batch_offset First batch index to process. + * @param batch_stride Distance between processed batch indices. + * + * @throws std::runtime_error if setup() has not been called. + */ + void solve(Vector rhs, int batch_offset = 0, int batch_stride = 1) override + { + if (!is_factorized_) { + throw std::runtime_error("Error: Matrix must be factorized before solving."); + } + + const int effective_batch_count = (this->batch_count_ - batch_offset + batch_stride - 1) / batch_stride; + + int matrix_dimension = this->matrix_dimension_; + int num_steps = num_steps_; + bool is_cyclic = this->is_cyclic_; + + Vector main_diagonal = this->main_diagonal_; + Vector sub_diagonal = this->sub_diagonal_; + Vector gamma = gamma_; + Vector k1_trajectory = k1_trajectory_; + Vector k2_trajectory = k2_trajectory_; + + if (matrix_dimension == 1) { + Kokkos::parallel_for( + "SolveTrivial", Kokkos::RangePolicy(0, effective_batch_count), + KOKKOS_LAMBDA(const int k) { + const int batch_idx = batch_stride * k + batch_offset; + rhs(batch_idx) /= main_diagonal(batch_idx); + }); + Kokkos::fence(); + return; + } + + using TeamPolicy = Kokkos::TeamPolicy; + using TeamMember = typename TeamPolicy::member_type; + + if (!is_cyclic) { + const std::size_t scratch_bytes = 2ull * static_cast(matrix_dimension) * sizeof(T); + + TeamPolicy policy(effective_batch_count, Kokkos::AUTO); + policy.set_scratch_size(0, Kokkos::PerTeam(static_cast(scratch_bytes))); + + Kokkos::parallel_for( + "SolveNonCyclicPCR", policy, KOKKOS_LAMBDA(const TeamMember& team_member) { + const int k = team_member.league_rank(); + const int batch_idx = batch_stride * k + batch_offset; + const int offset = batch_idx * matrix_dimension; + const int team_size = team_member.team_size(); + const int rank = team_member.team_rank(); + + T* scratch = static_cast(team_member.team_scratch(0).get_shmem(scratch_bytes)); + T* d[2] = {scratch, scratch + matrix_dimension}; + + int cur = 0; + for (int i = rank; i < matrix_dimension; i += team_size) { + d[cur][i] = rhs(offset + i); + } + team_member.team_barrier(); + + for (int step = 0; step < num_steps; step++) { + const int delta = 1 << step; + + for (int i = rank; i < matrix_dimension; i += team_size) { + int iLeft, iRight; + pcr_neighbors(i, delta, matrix_dimension, iLeft, iRight); + + const T k1_val = + k1_trajectory(static_cast(batch_idx) * num_steps * matrix_dimension + + static_cast(step) * matrix_dimension + i); + const T k2_val = + k2_trajectory(static_cast(batch_idx) * num_steps * matrix_dimension + + static_cast(step) * matrix_dimension + i); + + const int nxt = 1 - cur; + d[nxt][i] = d[cur][i] - d[cur][iLeft] * k1_val - d[cur][iRight] * k2_val; + } + + team_member.team_barrier(); + cur = 1 - cur; + } + + for (int i = rank; i < matrix_dimension; i += team_size) { + rhs(offset + i) = d[cur][i] / main_diagonal(offset + i); + } + }); + } + else { + // The cyclic solve simultaneously reduces the right-hand side and + // the auxiliary Sherman–Morrison vector. The original corner + // coefficient remains available in this->sub_diagonal_. + const std::size_t scratch_bytes = 4ull * static_cast(matrix_dimension) * sizeof(T); + + TeamPolicy policy(effective_batch_count, Kokkos::AUTO); + policy.set_scratch_size(0, Kokkos::PerTeam(static_cast(scratch_bytes))); + + Kokkos::parallel_for( + "SolveCyclicPCR", policy, KOKKOS_LAMBDA(const TeamMember& team_member) { + const int k = team_member.league_rank(); + const int batch_idx = batch_stride * k + batch_offset; + const int offset = batch_idx * matrix_dimension; + const int team_size = team_member.team_size(); + const int rank = team_member.team_rank(); + + T* scratch = static_cast(team_member.team_scratch(0).get_shmem(scratch_bytes)); + // Layout: [d_rhs(0) | d_buf(0) | d_rhs(1) | d_buf(1)]. + T* d_rhs[2] = {scratch, scratch + 2 * matrix_dimension}; + T* d_buf[2] = {scratch + matrix_dimension, scratch + 3 * matrix_dimension}; + + const T cyclic_corner_element = sub_diagonal(offset + matrix_dimension - 1); + + int cur = 0; + // Initial auxiliary vector for the Sherman–Morrison + // reconstruction. + for (int i = rank; i < matrix_dimension; i += team_size) { + d_rhs[cur][i] = rhs(offset + i); + if (i == 0) { + d_buf[cur][i] = gamma(batch_idx); + } + else if (i == matrix_dimension - 1) { + d_buf[cur][i] = cyclic_corner_element; + } + else { + d_buf[cur][i] = T(0); + } + } + + team_member.team_barrier(); + + for (int step = 0; step < num_steps; step++) { + const int delta = 1 << step; + + for (int i = rank; i < matrix_dimension; i += team_size) { + int iLeft, iRight; + pcr_neighbors(i, delta, matrix_dimension, iLeft, iRight); + + const T k1_val = + k1_trajectory(static_cast(batch_idx) * num_steps * matrix_dimension + + static_cast(step) * matrix_dimension + i); + const T k2_val = + k2_trajectory(static_cast(batch_idx) * num_steps * matrix_dimension + + static_cast(step) * matrix_dimension + i); + + const int nxt = 1 - cur; + d_rhs[nxt][i] = d_rhs[cur][i] - d_rhs[cur][iLeft] * k1_val - d_rhs[cur][iRight] * k2_val; + d_buf[nxt][i] = d_buf[cur][i] - d_buf[cur][iLeft] * k1_val - d_buf[cur][iRight] * k2_val; + } + + team_member.team_barrier(); + cur = 1 - cur; + } + + // The unused buffer stores the reduced solutions so all + // team members can access the entries required for the + // Sherman–Morrison reconstruction. + const int other = 1 - cur; + for (int i = rank; i < matrix_dimension; i += team_size) { + d_rhs[other][i] = d_rhs[cur][i] / main_diagonal(offset + i); + d_buf[other][i] = d_buf[cur][i] / main_diagonal(offset + i); + } + team_member.team_barrier(); + + const T dot_product_x_v = + d_rhs[other][0] + cyclic_corner_element / gamma(batch_idx) * d_rhs[other][matrix_dimension - 1]; + const T dot_product_u_v = + d_buf[other][0] + cyclic_corner_element / gamma(batch_idx) * d_buf[other][matrix_dimension - 1]; + const T factor = dot_product_x_v / (T(1) + dot_product_u_v); + + for (int i = rank; i < matrix_dimension; i += team_size) { + rhs(offset + i) = d_rhs[other][i] - factor * d_buf[other][i]; + } + }); + } + Kokkos::fence(); + } + + /** + * @brief Solves systems whose matrix has already been reduced to diagonal form. + * + * Each matrix entry is independent, so the operation is parallelized over + * both batch and equation indices. + * + * For cyclic systems, the first diagonal entry includes the corresponding + * Sherman–Morrison diagonal correction stored in gamma_. + * + * @param rhs Right-hand sides to solve. The vector is overwritten with the + * resulting solution. + * @param batch_offset First batch index to process. + * @param batch_stride Distance between processed batch indices. + * + * @throws std::runtime_error if setup() has not been called. + */ + void solve_diagonal(Vector rhs, int batch_offset = 0, int batch_stride = 1) override + { + if (!is_factorized_) { + throw std::runtime_error("Error: Matrix must be factorized before solving."); + } + + const int effective_batch_count = (this->batch_count_ - batch_offset + batch_stride - 1) / batch_stride; + + const int matrix_dimension = this->matrix_dimension_; + const bool is_cyclic = this->is_cyclic_; + Vector main_diagonal = this->main_diagonal_; + Vector gamma = gamma_; + + using MDPolicy = Kokkos::MDRangePolicy>; + MDPolicy policy({0, 0}, {effective_batch_count, matrix_dimension}); + + Kokkos::parallel_for( + "SolveDiagonal", policy, KOKKOS_LAMBDA(const int k, const int i) { + const int batch_idx = batch_stride * k + batch_offset; + const int offset = batch_idx * matrix_dimension; + + if (is_cyclic && i == 0) { + rhs(offset) /= main_diagonal(offset) + gamma(batch_idx); + } + else { + rhs(offset + i) /= main_diagonal(offset + i); + } + }); + + Kokkos::fence(); + } + +private: + int num_steps_; // Number of PCR reduction steps. + Vector k1_trajectory_; // Stored left reduction coefficients. + Vector k2_trajectory_; // Stored right reduction coefficients. + + Vector gamma_; + bool is_factorized_; +}; + +} // namespace gmgpolar diff --git a/include/LinearAlgebra/Solvers/tridiagonal_solver_thomas.h b/include/LinearAlgebra/Solvers/tridiagonal_solver_thomas.h new file mode 100644 index 00000000..8f1332d8 --- /dev/null +++ b/include/LinearAlgebra/Solvers/tridiagonal_solver_thomas.h @@ -0,0 +1,255 @@ +#pragma once + +#include + +#include "../../LinearAlgebra/Vector/vector.h" +#include "../../LinearAlgebra/Vector/vector_operations.h" + +namespace gmgpolar +{ + +template +class BatchedTridiagonalSolverThomas : public BatchedTridiagonalSolverBase +{ +public: + BatchedTridiagonalSolverThomas(int matrix_dimension, int batch_count, bool is_cyclic = true) + : BatchedTridiagonalSolverBase(matrix_dimension, batch_count, is_cyclic) + , buffer_("BatchedTridiagonalSolverThomas::buffer", is_cyclic ? matrix_dimension * batch_count : 0) + , gamma_("BatchedTridiagonalSolverThomas::gamma", is_cyclic ? batch_count : 0) + , is_factorized_(false) + { + assign(buffer_, T(0)); + assign(gamma_, T(0)); + } + + /* ---------------------------------------------- */ + /* Setup: Cholesky Decomposition: A = L * D * L^T */ + /* ---------------------------------------------- */ + // This step factorizes the tridiagonal matrix into lower triangular (L) and diagonal (D) matrices. + // For cyclic systems, it also applies the Shermann-Morrison adjustment to account for the cyclic connection. + + void setup() + { + // Create local copies for lambda capture + int matrix_dimension = this->matrix_dimension_; + Vector main_diagonal = this->main_diagonal_; + Vector sub_diagonal = this->sub_diagonal_; + Vector gamma = gamma_; + + if (!this->is_cyclic_) { + Kokkos::parallel_for( + "SetupNonCyclic", Kokkos::RangePolicy(0, this->batch_count_), + KOKKOS_LAMBDA(const int batch_idx) { + // ----------------------------------- // + // Obtain offset for the current batch // + int offset = batch_idx * matrix_dimension; + + // ---------------------- // + // Cholesky Decomposition // + for (int i = 1; i < matrix_dimension; i++) { + sub_diagonal(offset + i - 1) /= main_diagonal(offset + i - 1); + const T factor = sub_diagonal(offset + i - 1); + main_diagonal(offset + i) -= factor * factor * main_diagonal(offset + i - 1); + } + }); + } + else { + Kokkos::parallel_for( + "SetupCyclic", Kokkos::RangePolicy(0, this->batch_count_), + KOKKOS_LAMBDA(const int batch_idx) { + // ----------------------------------- // + // Obtain offset for the current batch // + int offset = batch_idx * matrix_dimension; + + // ------------------------------------------------- // + // Shermann-Morrison Adjustment // + // - Modify the first and last main diagonal element // + // - Compute and store gamma for later use // + // ------------------------------------------------- // + T cyclic_corner_element = sub_diagonal(offset + matrix_dimension - 1); + gamma(batch_idx) = -main_diagonal(offset + 0); + main_diagonal(offset + 0) -= gamma(batch_idx); + main_diagonal(offset + matrix_dimension - 1) -= + cyclic_corner_element * cyclic_corner_element / gamma(batch_idx); + + // ---------------------- // + // Cholesky Decomposition // + for (int i = 1; i < matrix_dimension; i++) { + sub_diagonal(offset + i - 1) /= main_diagonal(offset + i - 1); + const T factor = sub_diagonal(offset + i - 1); + main_diagonal(offset + i) -= factor * factor * main_diagonal(offset + i - 1); + } + }); + } + Kokkos::fence(); + is_factorized_ = true; + } + + /* ---------------------------------------- */ + /* Solve: Forward and Backward Substitution */ + /* ---------------------------------------- */ + // This step solves the system Ax = b using the factorized form of A. + // For cyclic systems, it also performs the Shermann-Morrison reconstruction to obtain the final solution. + + void solve(Vector rhs, int batch_offset = 0, int batch_stride = 1) + { + if (!is_factorized_) { + throw std::runtime_error("Error: Matrix must be factorized before solving."); + } + + // Compute the effective number of batches to solve + int effective_batch_count = (this->batch_count_ - batch_offset + batch_stride - 1) / batch_stride; + + // Create local copies for lambda capture + int matrix_dimension = this->matrix_dimension_; + Vector main_diagonal = this->main_diagonal_; + Vector sub_diagonal = this->sub_diagonal_; + Vector buffer = buffer_; + Vector gamma = gamma_; + + if (!this->is_cyclic_) { + Kokkos::parallel_for( + "SolveNonCyclic", Kokkos::RangePolicy(0, effective_batch_count), + KOKKOS_LAMBDA(const int k) { + // ----------------------------------- // + // Obtain offset for the current batch // + int batch_idx = batch_stride * k + batch_offset; + int offset = batch_idx * matrix_dimension; + + // -------------------- // + // Forward Substitution // + for (int i = 1; i < matrix_dimension; i++) { + rhs(offset + i) -= sub_diagonal(offset + i - 1) * rhs(offset + i - 1); + } + + // ---------------- // + // Diagonal Scaling // + for (int i = 0; i < matrix_dimension; i++) { + rhs(offset + i) /= main_diagonal(offset + i); + } + + // --------------------- // + // Backward Substitution // + for (int i = matrix_dimension - 2; i >= 0; i--) { + rhs(offset + i) -= sub_diagonal(offset + i) * rhs(offset + i + 1); + } + }); + } + else { + Kokkos::parallel_for( + "SolveCyclic", Kokkos::RangePolicy(0, effective_batch_count), + KOKKOS_LAMBDA(const int k) { + // ----------------------------------- // + // Obtain offset for the current batch // + int batch_idx = batch_stride * k + batch_offset; + int offset = batch_idx * matrix_dimension; + + // -------------------- // + // Forward Substitution // + T cyclic_corner_element = sub_diagonal(offset + matrix_dimension - 1); + buffer(offset + 0) = gamma(batch_idx); + for (int i = 1; i < matrix_dimension; i++) { + rhs(offset + i) -= sub_diagonal(offset + i - 1) * rhs(offset + i - 1); + if (i < matrix_dimension - 1) + buffer(offset + i) = 0.0 - sub_diagonal(offset + i - 1) * buffer(offset + i - 1); + else + buffer(offset + i) = + cyclic_corner_element - sub_diagonal(offset + i - 1) * buffer(offset + i - 1); + } + + // ---------------- // + // Diagonal Scaling // + for (int i = 0; i < matrix_dimension; i++) { + rhs(offset + i) /= main_diagonal(offset + i); + buffer(offset + i) /= main_diagonal(offset + i); + } + + // --------------------- // + // Backward Substitution // + for (int i = matrix_dimension - 2; i >= 0; i--) { + rhs(offset + i) -= sub_diagonal(offset + i) * rhs(offset + i + 1); + buffer(offset + i) -= sub_diagonal(offset + i) * buffer(offset + i + 1); + } + + // ------------------------------- // + // Shermann-Morrison Reonstruction // + const T dot_product_x_v = + rhs(offset + 0) + cyclic_corner_element / gamma(batch_idx) * rhs(offset + matrix_dimension - 1); + + const T dot_product_u_v = buffer(offset + 0) + cyclic_corner_element / gamma(batch_idx) * + buffer(offset + matrix_dimension - 1); + + const T factor = dot_product_x_v / (1.0 + dot_product_u_v); + + for (int i = 0; i < matrix_dimension; i++) { + rhs(offset + i) -= factor * buffer(offset + i); + } + }); + } + Kokkos::fence(); + } + + /* ---------------------------- */ + /* Solve: Diagonal Scaling Only */ + /* ---------------------------- */ + // This step performs only the diagonal scaling part of the solve process. + // It is useful when the matrix has a non-zero diagonal but zero off-diagonal entries. + // Note that .setup() modifies main_diagonal(0) in the cyclic case. + + void solve_diagonal(Vector rhs, int batch_offset = 0, int batch_stride = 1) + { + if (!is_factorized_) { + throw std::runtime_error("Error: Matrix must be factorized before solving."); + } + + // Compute the effective number of batches to solve + int effective_batch_count = (this->batch_count_ - batch_offset + batch_stride - 1) / batch_stride; + + // Create local copies for lambda capture + int matrix_dimension = this->matrix_dimension_; + Vector main_diagonal = this->main_diagonal_; + Vector gamma = gamma_; + + if (!this->is_cyclic_) { + Kokkos::parallel_for( + "SolveDiagonalNonCyclic", Kokkos::RangePolicy(0, effective_batch_count), + KOKKOS_LAMBDA(const int k) { + // ----------------------------------- // + // Obtain offset for the current batch // + int batch_idx = batch_stride * k + batch_offset; + int offset = batch_idx * matrix_dimension; + + // ---------------- // + // Diagonal Scaling // + for (int i = 0; i < matrix_dimension; i++) { + rhs(offset + i) /= main_diagonal(offset + i); + } + }); + } + else { + Kokkos::parallel_for( + "SolveDiagonalCyclic", Kokkos::RangePolicy(0, effective_batch_count), + KOKKOS_LAMBDA(const int k) { + // ----------------------------------- // + // Obtain offset for the current batch // + int batch_idx = batch_stride * k + batch_offset; + int offset = batch_idx * matrix_dimension; + + // ---------------- // + // Diagonal Scaling // + rhs(offset + 0) /= main_diagonal(offset + 0) + gamma(batch_idx); + for (int i = 1; i < matrix_dimension; i++) { + rhs(offset + i) /= main_diagonal(offset + i); + } + }); + } + Kokkos::fence(); + } + +private: + Vector buffer_; + Vector gamma_; + + bool is_factorized_; +}; +} // namespace gmgpolar diff --git a/tests/LinearAlgebra/Solvers/tridiagonal_solver.cpp b/tests/LinearAlgebra/Solvers/tridiagonal_solver.cpp index 3a3e4330..3443fdac 100644 --- a/tests/LinearAlgebra/Solvers/tridiagonal_solver.cpp +++ b/tests/LinearAlgebra/Solvers/tridiagonal_solver.cpp @@ -2,19 +2,52 @@ #include #include #include +#include +#include #include + #include using namespace gmgpolar; // clang-format off + +// ----------------------------------------------------------------------------------------------- +// Every test in this file is written as a function template on the concrete solver type and is +// instantiated for all four batched tridiagonal solver backends: Thomas, CR, PCR and CRPCR. +// +// Previously these tests only used `BatchedTridiagonalSolver`, a type alias that picks a +// SINGLE one of the four implementations depending on the active Kokkos backend (Thomas on +// host-only builds, CRPCR when CUDA/HIP/SYCL is enabled). That means CR and PCR were never +// exercised by this file at all, and depending on the build, either Thomas or CRPCR was silently +// skipped too. +// +// Test bodies are still free functions (not TEST/TYPED_TEST member functions) for the same reason +// as before: CUDA extended device lambdas cannot be defined inside a function with internal +// linkage, so the Kokkos::parallel_for bodies need to live in ordinary, externally-linked +// functions. The INSTANTIATE_FOR_ALL_SOLVERS macro below generates the four TEST(...) cases (one +// per solver backend) that each call the templated test function with the corresponding solver +// type. +// ----------------------------------------------------------------------------------------------- + +#define INSTANTIATE_FOR_ALL_SOLVERS(TESTNAME) \ + TEST(BatchedTridiagonalSolvers_Thomas, TESTNAME) { test_##TESTNAME>(); } \ + TEST(BatchedTridiagonalSolvers_CR, TESTNAME) { test_##TESTNAME>(); } \ + TEST(BatchedTridiagonalSolvers_PCR, TESTNAME) { test_##TESTNAME>(); } \ + TEST(BatchedTridiagonalSolvers_CRPCR, TESTNAME) { test_##TESTNAME>(); } + +// ================================================================================================= +// Hand-derived 4x4 systems (unchanged from before, just templated on SolverType) +// ================================================================================================= + +template void test_non_cyclic_tridiagonal_n_4() { int batch_count = 4; int matrix_dimension = 4; bool is_cyclic = false; - BatchedTridiagonalSolver solver(matrix_dimension, batch_count, is_cyclic); + SolverType solver(matrix_dimension, batch_count, is_cyclic); // System 1: {{2, 1, 0,0},{1,4,2,0},{0,2,6,3},{0,0,3,8}} * {{a},{b},{c},{d}} = {{1},{2},{3},{4}} // a = 70/209, b = 69/209, c = 36/209, d = 91/209 @@ -29,6 +62,8 @@ void test_non_cyclic_tridiagonal_n_4() "Test", 1, KOKKOS_LAMBDA(const int) { + + solver.set_main_diagonal(0,0, 2.0); solver.set_sub_diagonal(0,0, 1.0); solver.set_main_diagonal(0,1, 4.0); solver.set_sub_diagonal(0,1, 2.0); solver.set_main_diagonal(0,2, 6.0); solver.set_sub_diagonal(0,2, 3.0); @@ -95,19 +130,16 @@ void test_non_cyclic_tridiagonal_n_4() EXPECT_NEAR(h_rhs(14), 267.0/710.0, tol); EXPECT_NEAR(h_rhs(15), 379.0/710.0, tol); } -TEST(BatchedTridiagonalSolvers, non_cyclic_tridiagonal_n_4) -{ - // Call a function named function due to cuda restriction - test_non_cyclic_tridiagonal_n_4(); -} +INSTANTIATE_FOR_ALL_SOLVERS(non_cyclic_tridiagonal_n_4) +template void test_cyclic_tridiagonal_n_4() { int batch_count = 4; int matrix_dimension = 4; bool is_cyclic = true; - BatchedTridiagonalSolver solver(matrix_dimension, batch_count, is_cyclic); + SolverType solver(matrix_dimension, batch_count, is_cyclic); // System 1: {{2, 1, 0,-1},{1,4,2,0},{0,2,6,3},{-1,0,3,8}} * {{a},{b},{c},{d}} = {{1},{2},{3},{4}} // a = 42/67, b = 18/67, c = 10/67, d = 35/67 @@ -188,17 +220,15 @@ void test_cyclic_tridiagonal_n_4() EXPECT_NEAR(h_rhs(14), 14.0/81.0, tol); EXPECT_NEAR(h_rhs(15), 97.0/81.0, tol); } -TEST(BatchedTridiagonalSolvers, cyclic_tridiagonal_n_4) { - // Call a function named function due to cuda restriction - test_cyclic_tridiagonal_n_4(); -} +INSTANTIATE_FOR_ALL_SOLVERS(cyclic_tridiagonal_n_4) +template void test_non_cyclic_diagonal_n_4() { int batch_count = 4; int matrix_dimension = 4; bool is_cyclic = false; - BatchedTridiagonalSolver solver(matrix_dimension, batch_count, is_cyclic); + SolverType solver(matrix_dimension, batch_count, is_cyclic); Kokkos::parallel_for( "Test", @@ -270,17 +300,15 @@ void test_non_cyclic_diagonal_n_4() { EXPECT_NEAR(h_rhs(14), 6.0/9.0, tol); EXPECT_NEAR(h_rhs(15), 7.0/11.0, tol); } -TEST(BatchedTridiagonalSolvers, non_cyclic_diagonal_n_4) { - // Call a function named function due to cuda restriction - test_non_cyclic_diagonal_n_4(); -} +INSTANTIATE_FOR_ALL_SOLVERS(non_cyclic_diagonal_n_4) +template void test_cyclic_diagonal_n_4() { int batch_count = 4; int matrix_dimension = 4; bool is_cyclic = true; - BatchedTridiagonalSolver solver(matrix_dimension, batch_count, is_cyclic); + SolverType solver(matrix_dimension, batch_count, is_cyclic); Kokkos::parallel_for( "Test", @@ -352,7 +380,679 @@ void test_cyclic_diagonal_n_4() { EXPECT_NEAR(h_rhs(14), 6.0/9.0, tol); EXPECT_NEAR(h_rhs(15), 7.0/11.0, tol); } -TEST(BatchedTridiagonalSolvers, cyclic_diagonal_n_4) { - // Call a function named function due to cuda restriction - test_cyclic_diagonal_n_4(); +INSTANTIATE_FOR_ALL_SOLVERS(cyclic_diagonal_n_4) + + +// ----------------------------------------------------------------------------------------------- +// Additional edge-case coverage (templated on SolverType, unchanged in spirit from before). +// +// The four tests above only ever exercise matrix_dimension_ == 4 (a clean power of two, well +// above every boundary-clamping/degenerate special case) with batch_count_ == matrix_dimension_ +// and a perfectly even 50/50 stride split. The tests below additionally touch: +// - matrix_dimension_ == 1 (the early-return branch in setup(), and the RangePolicy-only branch +// in solve()) +// - matrix_dimension_ == 2 (every PCR/CR step clamps neighbors to the opposite boundary) +// - non-power-of-two matrix_dimension_ (num_steps_ = ceil(log2(n)) taking a non-exact value) +// - batch_count_ == 1 (TeamPolicy league_size == 1) +// - batch_count_ > matrix_dimension_ (the inverse of the "few, long lines" regime the file's own +// header comments assume) +// - matrix_dimension_ large enough that Kokkos::AUTO's chosen team_size is very likely < n, +// forcing strided team loops to actually iterate more than once per thread +// - batch_offset_/batch_stride_ combinations that don't split batch_count_ evenly +// - calling solve()/solve_diagonal() with their default arguments at all +// - matrix_dimension_ in the thousands, well beyond anything reachable via dense O(n^3) +// Gaussian-elimination verification +// +// Since hand-deriving exact fractions stops being practical past n=4, these tests check against +// independently-implemented reference solvers instead. They share no code with any of the four +// solver classes under test - only the mathematics of solving the actual linear system. +// ----------------------------------------------------------------------------------------------- + +// Deterministic (not random) but distinct-per-(batch,index) diagonally dominant system generator, +// shared between "build the solver's input" and "build the reference's input" below, so both +// sides are guaranteed to describe the same matrix without needing a host<->device data transfer +// of arbitrary-sized arrays into a KOKKOS_LAMBDA. +KOKKOS_INLINE_FUNCTION double sysDiag(int b, int i) +{ + return 8.0 + 2.0 * i + 0.5 * b; // baseline large enough to keep the corner + off-diagonals + // diagonally dominant up to the largest n/batch used below +} +KOKKOS_INLINE_FUNCTION double sysSub(int b, int i) +{ + return 1.0 + 0.1 * i; +} +KOKKOS_INLINE_FUNCTION double sysCorner(int b) +{ + return 0.7 + 0.05 * b; +} +KOKKOS_INLINE_FUNCTION double sysRhs(int b, int i) +{ + return 1.0 + 0.3 * i + 0.2 * b; +} + +// Fully independent reference for small/moderate n: build the dense n x n matrix explicitly +// (tridiagonal, plus the wraparound corner entries when is_cyclic) and solve it with plain +// partial-pivot Gaussian elimination. O(n^3) is irrelevant at the sizes used for these tests +// (up to ~1024) but far too slow for the n=4000-scale tests further below. +static void denseReferenceSolve(int n, const std::vector& diag, const std::vector& subdiag, + double corner, bool is_cyclic, std::vector& rhs) +{ + + std::vector> A(n, std::vector(n, 0.0)); + for (int i = 0; i < n; i++) { + A[i][i] = diag[i]; + } + for (int i = 0; i < n - 1; i++) { + A[i][i + 1] = subdiag[i]; + A[i + 1][i] = subdiag[i]; + } + if (is_cyclic && n > 1) { + A[0][n - 1] += corner; + A[n - 1][0] += corner; + } + + for (int col = 0; col < n; col++) { + int pivot = col; + for (int row = col + 1; row < n; row++) { + if (std::fabs(A[row][col]) > std::fabs(A[pivot][col])) + pivot = row; + } + std::swap(A[col], A[pivot]); + std::swap(rhs[col], rhs[pivot]); + + for (int row = col + 1; row < n; row++) { + double factor = A[row][col] / A[col][col]; + for (int c = col; c < n; c++) { + A[row][c] -= factor * A[col][c]; + } + rhs[row] -= factor * rhs[col]; + } + } + + for (int row = n - 1; row >= 0; row--) { + double sum = rhs[row]; + for (int c = row + 1; c < n; c++) { + sum -= A[row][c] * rhs[c]; + } + rhs[row] = sum / A[row][row]; + } +} + +// Independent O(n) banded reference solver, used only for the very-large-matrix_dimension_ tests +// where the O(n^3) dense Gaussian elimination above would be impractically slow (n=4000 => ~6.4e10 +// flops per system). Any exact direct solve of a tridiagonal system is necessarily some variant of +// forward-elimination/back-substitution - there is no way to be algorithmically "independent" of +// that shape at O(n) - but this implementation is written from scratch here and does not call, +// share code with, or share any intermediate representation with any of the four solver classes +// under test (Thomas/CR/PCR/CRPCR), so it still independently catches bugs in each of them. +static void bandedReferenceSolveNonCyclic(int n, const std::vector& diag, const std::vector& subdiag, + std::vector& rhs) +{ + std::vector c_prime(n, 0.0), d_prime(n, 0.0); + c_prime[0] = (n > 1) ? subdiag[0] / diag[0] : 0.0; + d_prime[0] = rhs[0] / diag[0]; + for (int i = 1; i < n; i++) { + double denom = diag[i] - subdiag[i - 1] * c_prime[i - 1]; + c_prime[i] = (i < n - 1) ? subdiag[i] / denom : 0.0; + d_prime[i] = (rhs[i] - subdiag[i - 1] * d_prime[i - 1]) / denom; + } + rhs[n - 1] = d_prime[n - 1]; + for (int i = n - 2; i >= 0; i--) { + rhs[i] = d_prime[i] - c_prime[i] * rhs[i + 1]; + } +} + +// Cyclic counterpart via the standard Sherman-Morrison technique for bordered tridiagonal systems +// (turns one cyclic solve into two plain tridiagonal solves plus a rank-1 correction). n == 2 is +// special-cased because there the "corner" position (0,1)/(1,0) coincides with the regular +// off-diagonal position rather than being a separate matrix entry (matching how denseReferenceSolve +// and the solver classes under test build their n==2 cyclic matrix). +static void bandedReferenceSolveCyclic(int n, const std::vector& diag, const std::vector& subdiag, + double corner, std::vector& rhs) +{ + if (n == 1) { + rhs[0] = rhs[0] / (diag[0] + 2.0 * corner); + return; + } + if (n == 2) { + double a00 = diag[0]; + double aoff = subdiag[0] + corner; + double a11 = diag[1]; + double det = a00 * a11 - aoff * aoff; + double r0 = rhs[0], r1 = rhs[1]; + rhs[0] = (r0 * a11 - aoff * r1) / det; + rhs[1] = (a00 * r1 - aoff * r0) / det; + return; + } + + double alpha = corner; + double beta = corner; + double gamma = -diag[0]; + + std::vector bb = diag; + bb[0] -= gamma; + bb[n - 1] -= alpha * beta / gamma; + + std::vector u(n, 0.0); + u[0] = gamma; + u[n - 1] = alpha; + + std::vector x = rhs; + bandedReferenceSolveNonCyclic(n, bb, subdiag, x); + + std::vector z = u; + bandedReferenceSolveNonCyclic(n, bb, subdiag, z); + + double fact = (x[0] + beta * x[n - 1] / gamma) / (1.0 + z[0] + beta * z[n - 1] / gamma); + + for (int i = 0; i < n; i++) { + rhs[i] = x[i] - fact * z[i]; + } +} + +// Builds a solver instance sized (matrix_dimension, batch_count, is_cyclic), fills it with the +// deterministic sysDiag/sysSub/sysCorner/sysRhs data, calls setup() once then solve() with the +// given (batch_offset, batch_stride), and checks every system that offset/stride combination +// actually covers against denseReferenceSolve. Does NOT assume batch_offset/stride cover the +// whole batch - callers needing full coverage make multiple calls against the same solver (see +// test_uneven_batch_stride_split below) or pass offset=0/stride=1. Intended for small/moderate n +// (dense O(n^3) reference) - use runHugeCorrectnessCheck for n in the thousands. +template +static void runCorrectnessCheck(int matrix_dimension, int batch_count, bool is_cyclic, int batch_offset, + int batch_stride) +{ + SolverType solver(matrix_dimension, batch_count, is_cyclic); + + Kokkos::parallel_for( + "Fill", 1, KOKKOS_LAMBDA(const int) { + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + solver.set_main_diagonal(b, i, sysDiag(b, i)); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + solver.set_sub_diagonal(b, i, sysSub(b, i)); + } + if (is_cyclic && matrix_dimension > 1) { + solver.set_cyclic_corner(b, sysCorner(b)); + } + } + }); + + HostVector h_rhs("h_rhs", matrix_dimension * batch_count); + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + h_rhs(b * matrix_dimension + i) = sysRhs(b, i); + } + } + + auto rhs = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), h_rhs); + solver.setup(); + solver.solve(rhs, batch_offset, batch_stride); + Kokkos::deep_copy(h_rhs, rhs); + + const double tol = 1e-9; + for (int k = 0;; k++) { + int batch_idx = batch_stride * k + batch_offset; + if (batch_idx >= batch_count) + break; + + std::vector diag(matrix_dimension); + std::vector subdiag(matrix_dimension > 1 ? matrix_dimension - 1 : 0); + std::vector rhs_ref(matrix_dimension); + for (int i = 0; i < matrix_dimension; i++) { + diag[i] = sysDiag(batch_idx, i); + rhs_ref[i] = sysRhs(batch_idx, i); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + subdiag[i] = sysSub(batch_idx, i); + } + double corner = sysCorner(batch_idx); + + denseReferenceSolve(matrix_dimension, diag, subdiag, corner, is_cyclic, rhs_ref); + + for (int i = 0; i < matrix_dimension; i++) { + EXPECT_NEAR(h_rhs(batch_idx * matrix_dimension + i), rhs_ref[i], tol) + << "n=" << matrix_dimension << " batch=" << batch_idx << " index=" << i; + } + } +} + +// Same idea as runCorrectnessCheck but for matrix_dimension_ in the thousands: uses the O(n) +// bandedReferenceSolve{Non}Cyclic reference instead of the O(n^3) dense one. +template +static void runHugeCorrectnessCheck(int matrix_dimension, int batch_count, bool is_cyclic, int batch_offset, + int batch_stride) +{ + SolverType solver(matrix_dimension, batch_count, is_cyclic); + + Kokkos::parallel_for( + "FillHuge", 1, KOKKOS_LAMBDA(const int) { + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + solver.set_main_diagonal(b, i, sysDiag(b, i)); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + solver.set_sub_diagonal(b, i, sysSub(b, i)); + } + if (is_cyclic && matrix_dimension > 1) { + solver.set_cyclic_corner(b, sysCorner(b)); + } + } + }); + + HostVector h_rhs("h_rhs", matrix_dimension * batch_count); + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + h_rhs(b * matrix_dimension + i) = sysRhs(b, i); + } + } + + auto rhs = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), h_rhs); + solver.setup(); + solver.solve(rhs, batch_offset, batch_stride); + Kokkos::deep_copy(h_rhs, rhs); + + const double tol = 1e-6; + for (int k = 0;; k++) { + int batch_idx = batch_stride * k + batch_offset; + if (batch_idx >= batch_count) + break; + + std::vector diag(matrix_dimension); + std::vector subdiag(matrix_dimension > 1 ? matrix_dimension - 1 : 0); + std::vector rhs_ref(matrix_dimension); + for (int i = 0; i < matrix_dimension; i++) { + diag[i] = sysDiag(batch_idx, i); + rhs_ref[i] = sysRhs(batch_idx, i); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + subdiag[i] = sysSub(batch_idx, i); + } + + if (is_cyclic) { + bandedReferenceSolveCyclic(matrix_dimension, diag, subdiag, sysCorner(batch_idx), rhs_ref); + } else { + bandedReferenceSolveNonCyclic(matrix_dimension, diag, subdiag, rhs_ref); + } + + for (int i = 0; i < matrix_dimension; i++) { + EXPECT_NEAR(h_rhs(batch_idx * matrix_dimension + i), rhs_ref[i], tol) + << "n=" << matrix_dimension << " batch=" << batch_idx << " index=" << i; + } + } +} + +// Diagonal-only counterpart: solve_diagonal() deliberately ignores off-diagonal coupling, so its +// only well-defined expected answer is plain rhs[i] / diag[i] against the ORIGINAL diagonal - the +// solver is filled with zero off-diagonal entries (including a zero corner even when is_cyclic) so +// that expectation actually holds. See setup()'s Sherman-Morrison-Woodbury b[0] adjustment: for a +// genuinely diagonal cyclic system (corner == 0) the +gamma/-gamma correction solve_diagonal() +// applies cancels back to the original diagonal exactly - a nonzero corner here would not be a +// "diagonal matrix" in the first place, so it's intentionally not tested through this path. +template +static void runDiagonalCorrectnessCheck(int matrix_dimension, int batch_count, bool is_cyclic, int batch_offset, + int batch_stride) +{ + SolverType solver(matrix_dimension, batch_count, is_cyclic); + + Kokkos::parallel_for( + "FillDiag", 1, KOKKOS_LAMBDA(const int) { + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + solver.set_main_diagonal(b, i, sysDiag(b, i)); + } + } + }); + + HostVector h_rhs("h_rhs", matrix_dimension * batch_count); + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + h_rhs(b * matrix_dimension + i) = sysRhs(b, i); + } + } + + auto rhs = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), h_rhs); + solver.setup(); + solver.solve_diagonal(rhs, batch_offset, batch_stride); + Kokkos::deep_copy(h_rhs, rhs); + + const double tol = 1e-12; + for (int k = 0;; k++) { + int batch_idx = batch_stride * k + batch_offset; + if (batch_idx >= batch_count) + break; + for (int i = 0; i < matrix_dimension; i++) { + double expected = sysRhs(batch_idx, i) / sysDiag(batch_idx, i); + EXPECT_NEAR(h_rhs(batch_idx * matrix_dimension + i), expected, tol) + << "n=" << matrix_dimension << " batch=" << batch_idx << " index=" << i; + } + } +} + +// --- matrix_dimension_ == 1: the degenerate early-return in setup() and the RangePolicy-only +// branch in solve() - not entered by the hand-derived n=4 tests. Also checks solve() and +// solve_diagonal() agree with each other here, since for n=1 there's no off-diagonal coupling to +// distinguish them. +template +void test_trivial_matrix_dimension_1() +{ + for (bool is_cyclic : {false, true}) { + runCorrectnessCheck(/*n=*/1, /*batch_count=*/3, is_cyclic, /*offset=*/0, /*stride=*/1); + runDiagonalCorrectnessCheck(/*n=*/1, /*batch_count=*/3, is_cyclic, /*offset=*/0, /*stride=*/1); + } +} +INSTANTIATE_FOR_ALL_SOLVERS(trivial_matrix_dimension_1) + +// --- matrix_dimension_ == 2: smallest nontrivial size - every PCR/CR step's iLeft/iRight clamps +// to the opposite boundary immediately (delta=1 already reaches both ends of a length-2 system). +template +void test_matrix_dimension_2() +{ + for (bool is_cyclic : {false, true}) { + runCorrectnessCheck(/*n=*/2, /*batch_count=*/3, is_cyclic, /*offset=*/0, /*stride=*/1); + } +} +INSTANTIATE_FOR_ALL_SOLVERS(matrix_dimension_2) + +// --- Non-power-of-two matrix_dimension_: num_steps_ = ceil(log2(n)) is not exact, so the last +// reduction step operates on a system that's "overshot" past what n strictly requires. Covers odd, +// even, and values straddling a power-of-two boundary on both sides (3, 5, 6, 7 bracket 4 and 8). +template +void test_non_power_of_two_dimensions() +{ + for (int n : {3, 5, 6, 7}) { + for (bool is_cyclic : {false, true}) { + runCorrectnessCheck(n, /*batch_count=*/3, is_cyclic, /*offset=*/0, /*stride=*/1); + } + } +} +INSTANTIATE_FOR_ALL_SOLVERS(non_power_of_two_dimensions) + +// --- batch_count_ == 1: TeamPolicy league_size == 1, and batch_count_ > matrix_dimension_: the +// inverse of the "few, long lines" regime the file's parallelization-model comments assume. +template +void test_batch_count_extremes() +{ + for (bool is_cyclic : {false, true}) { + runCorrectnessCheck(/*n=*/5, /*batch_count=*/1, is_cyclic, /*offset=*/0, /*stride=*/1); + runCorrectnessCheck(/*n=*/2, /*batch_count=*/10, is_cyclic, /*offset=*/0, /*stride=*/1); + } +} +INSTANTIATE_FOR_ALL_SOLVERS(batch_count_extremes) + +// --- matrix_dimension_ large enough that Kokkos::AUTO's chosen team_size is plausibly < n on at +// least some backends, forcing strided team loops to actually run more than one iteration per +// thread - the hand-derived tests all use n=4, small enough that most backends would just grant +// team_size >= n and never exercise that loop's stride > 0 case at all. +template +void test_large_dimension_strided_team_loop() +{ + for (bool is_cyclic : {false, true}) { + runCorrectnessCheck(/*n=*/33, /*batch_count=*/2, is_cyclic, /*offset=*/0, /*stride=*/1); + } + runDiagonalCorrectnessCheck(/*n=*/33, /*batch_count=*/2, /*is_cyclic=*/true, /*offset=*/0, + /*stride=*/1); +} +INSTANTIATE_FOR_ALL_SOLVERS(large_dimension_strided_team_loop) + +// --- Uneven batch_offset_/batch_stride_ split: the hand-derived tests split batch_count_ evenly +// in half (offset 0/1, stride 2, batch_count_ divisible by 2). Here batch_count_ = 5 is not +// divisible by stride = 3, so effective_batch_count_'s ceiling-division arithmetic produces a +// genuinely uneven split (2, 2, 1) across three solve() calls sharing ONE setup() - also checks +// that calling solve() more than twice against the same stored trajectory still works. +template +void test_uneven_batch_stride_split() +{ + const int matrix_dimension = 4; + const int batch_count = 5; + + for (bool is_cyclic : {false, true}) { + SolverType solver(matrix_dimension, batch_count, is_cyclic); + + Kokkos::parallel_for( + "Fill", 1, KOKKOS_LAMBDA(const int) { + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + solver.set_main_diagonal(b, i, sysDiag(b, i)); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + solver.set_sub_diagonal(b, i, sysSub(b, i)); + } + if (is_cyclic) { + solver.set_cyclic_corner(b, sysCorner(b)); + } + } + }); + + HostVector h_rhs("h_rhs", matrix_dimension * batch_count); + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + h_rhs(b * matrix_dimension + i) = sysRhs(b, i); + } + } + + auto rhs = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), h_rhs); + solver.setup(); + + const int stride = 3; + solver.solve(rhs, /*offset=*/0, stride); // covers batch 0, 3 + solver.solve(rhs, /*offset=*/1, stride); // covers batch 1, 4 + solver.solve(rhs, /*offset=*/2, stride); // covers batch 2 + + Kokkos::deep_copy(h_rhs, rhs); + + const double tol = 1e-9; + for (int b = 0; b < batch_count; b++) { + std::vector diag(matrix_dimension), subdiag(matrix_dimension - 1), rhs_ref(matrix_dimension); + for (int i = 0; i < matrix_dimension; i++) { + diag[i] = sysDiag(b, i); + rhs_ref[i] = sysRhs(b, i); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + subdiag[i] = sysSub(b, i); + } + denseReferenceSolve(matrix_dimension, diag, subdiag, sysCorner(b), is_cyclic, rhs_ref); + + for (int i = 0; i < matrix_dimension; i++) { + EXPECT_NEAR(h_rhs(b * matrix_dimension + i), rhs_ref[i], tol) << "batch " << b << " index " << i; + } + } + } +} +INSTANTIATE_FOR_ALL_SOLVERS(uneven_batch_stride_split) + +// --- Default batch_offset_/batch_stride_ arguments: checks solve()/solve_diagonal() called with +// no arguments at all (full batch in a single call) still compile and behave as offset=0, stride=1. +template +void test_default_batch_arguments() +{ + const int matrix_dimension = 4; + const int batch_count = 3; + + for (bool is_cyclic : {false, true}) { + SolverType solver(matrix_dimension, batch_count, is_cyclic); + + Kokkos::parallel_for( + "Fill", 1, KOKKOS_LAMBDA(const int) { + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + solver.set_main_diagonal(b, i, sysDiag(b, i)); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + solver.set_sub_diagonal(b, i, sysSub(b, i)); + } + if (is_cyclic) { + solver.set_cyclic_corner(b, sysCorner(b)); + } + } + }); + + HostVector h_rhs("h_rhs", matrix_dimension * batch_count); + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + h_rhs(b * matrix_dimension + i) = sysRhs(b, i); + } + } + + auto rhs = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), h_rhs); + solver.setup(); + solver.solve(rhs); // no offset/stride passed + + Kokkos::deep_copy(h_rhs, rhs); + + const double tol = 1e-9; + for (int b = 0; b < batch_count; b++) { + std::vector diag(matrix_dimension), subdiag(matrix_dimension - 1), rhs_ref(matrix_dimension); + for (int i = 0; i < matrix_dimension; i++) { + diag[i] = sysDiag(b, i); + rhs_ref[i] = sysRhs(b, i); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + subdiag[i] = sysSub(b, i); + } + denseReferenceSolve(matrix_dimension, diag, subdiag, sysCorner(b), is_cyclic, rhs_ref); + + for (int i = 0; i < matrix_dimension; i++) { + EXPECT_NEAR(h_rhs(b * matrix_dimension + i), rhs_ref[i], tol) << "batch " << b << " index " << i; + } + } + } +} +INSTANTIATE_FOR_ALL_SOLVERS(default_batch_arguments) + +// --- Large matrix_dimension_ (still within reach of the O(n^3) dense reference): exercises the +// reduction recursion over many more steps than the hand-derived tests, and stresses the strided +// team loops with a dimension unlikely to fit inside a single team's thread count on any backend. +// Covers both a non-power-of-two size (513, just past 2^9) and an exact power-of-two size (1024). +template +void test_very_large_matrix_dimension_non_cyclic() +{ + runCorrectnessCheck(/*n=*/513, /*batch_count=*/2, /*is_cyclic=*/false, /*offset=*/0, /*stride=*/1); + runCorrectnessCheck(/*n=*/1024, /*batch_count=*/2, /*is_cyclic=*/false, /*offset=*/0, /*stride=*/1); +} +INSTANTIATE_FOR_ALL_SOLVERS(very_large_matrix_dimension_non_cyclic) + +template +void test_very_large_matrix_dimension_cyclic() +{ + runCorrectnessCheck(/*n=*/513, /*batch_count=*/2, /*is_cyclic=*/true, /*offset=*/0, /*stride=*/1); + runCorrectnessCheck(/*n=*/1024, /*batch_count=*/2, /*is_cyclic=*/true, /*offset=*/0, /*stride=*/1); +} +INSTANTIATE_FOR_ALL_SOLVERS(very_large_matrix_dimension_cyclic) + +// --- Large matrix_dimension_ combined with a non-trivial batch_offset_/batch_stride_ split: +// checks that the batch-selection arithmetic and the large-n recursion compose correctly together, +// rather than each only being tested in isolation. batch_count_ = 5 with stride = 2 covers batches +// {0, 2, 4} on one call and {1, 3} on the other, sharing a single setup() at large n. +template +void test_very_large_matrix_dimension_with_batch_stride_non_cyclic() +{ + const int matrix_dimension = 777; + const int batch_count = 5; + const bool is_cyclic = false; + + SolverType solver(matrix_dimension, batch_count, is_cyclic); + + Kokkos::parallel_for( + "Fill", 1, KOKKOS_LAMBDA(const int) { + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + solver.set_main_diagonal(b, i, sysDiag(b, i)); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + solver.set_sub_diagonal(b, i, sysSub(b, i)); + } + } + }); + + HostVector h_rhs("h_rhs", matrix_dimension * batch_count); + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + h_rhs(b * matrix_dimension + i) = sysRhs(b, i); + } + } + + auto rhs = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), h_rhs); + solver.setup(); + + const int stride = 2; + solver.solve(rhs, /*offset=*/0, stride); // covers batch 0, 2, 4 + solver.solve(rhs, /*offset=*/1, stride); // covers batch 1, 3 + + Kokkos::deep_copy(h_rhs, rhs); + + const double tol = 1e-8; + for (int b = 0; b < batch_count; b++) { + std::vector diag(matrix_dimension), subdiag(matrix_dimension - 1), rhs_ref(matrix_dimension); + for (int i = 0; i < matrix_dimension; i++) { + diag[i] = sysDiag(b, i); + rhs_ref[i] = sysRhs(b, i); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + subdiag[i] = sysSub(b, i); + } + denseReferenceSolve(matrix_dimension, diag, subdiag, sysCorner(b), is_cyclic, rhs_ref); + + for (int i = 0; i < matrix_dimension; i++) { + EXPECT_NEAR(h_rhs(b * matrix_dimension + i), rhs_ref[i], tol) << "batch " << b << " index " << i; + } + } +} +INSTANTIATE_FOR_ALL_SOLVERS(very_large_matrix_dimension_with_batch_stride_non_cyclic) + +template +void test_very_large_matrix_dimension_with_batch_stride_cyclic() +{ + const int matrix_dimension = 777; + const int batch_count = 5; + const bool is_cyclic = true; + + SolverType solver(matrix_dimension, batch_count, is_cyclic); + + Kokkos::parallel_for( + "Fill", 1, KOKKOS_LAMBDA(const int) { + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + solver.set_main_diagonal(b, i, sysDiag(b, i)); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + solver.set_sub_diagonal(b, i, sysSub(b, i)); + } + solver.set_cyclic_corner(b, sysCorner(b)); + } + }); + + HostVector h_rhs("h_rhs", matrix_dimension * batch_count); + for (int b = 0; b < batch_count; b++) { + for (int i = 0; i < matrix_dimension; i++) { + h_rhs(b * matrix_dimension + i) = sysRhs(b, i); + } + } + + auto rhs = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), h_rhs); + solver.setup(); + + const int stride = 2; + solver.solve(rhs, /*offset=*/0, stride); // covers batch 0, 2, 4 + solver.solve(rhs, /*offset=*/1, stride); // covers batch 1, 3 + + Kokkos::deep_copy(h_rhs, rhs); + + const double tol = 1e-8; + for (int b = 0; b < batch_count; b++) { + std::vector diag(matrix_dimension), subdiag(matrix_dimension - 1), rhs_ref(matrix_dimension); + for (int i = 0; i < matrix_dimension; i++) { + diag[i] = sysDiag(b, i); + rhs_ref[i] = sysRhs(b, i); + } + for (int i = 0; i < matrix_dimension - 1; i++) { + subdiag[i] = sysSub(b, i); + } + denseReferenceSolve(matrix_dimension, diag, subdiag, sysCorner(b), is_cyclic, rhs_ref); + + for (int i = 0; i < matrix_dimension; i++) { + EXPECT_NEAR(h_rhs(b * matrix_dimension + i), rhs_ref[i], tol) << "batch " << b << " index " << i; + } + } } +INSTANTIATE_FOR_ALL_SOLVERS(very_large_matrix_dimension_with_batch_stride_cyclic)