Batched Linear Algebra

Parthenon provides a small set of dense linear algebra routines in src/batched_linear_algebra/ (namespace parthenon::batched_linear_algebra). They are meant for many small, independent problems, for example one matrix per cell or per meshblock. Each routine can be called from host code, by a single thread in a flat kernel, or by a Kokkos team in a hierarchical kernel (see Ways to call the routines).

Note

These routines are designed for small matrices (up to a few tens of rows and columns). They are not a replacement for a distributed or vendor BLAS/LAPACK for large problems.

Available decompositions

All routines work in place and overwrite the input matrix. Optional output factors are passed as pointers and may be nullptr if they are not needed.

QRDecomposition::execute(tm, pA, pQ, scratch): Householder QR of an \(m \times n\) matrix with \(m \geq n\). On exit *pA holds the upper-trapezoidal \(R\), and all entries below the diagonal are exactly zero. *pQ may be a full (\(m \times m\)) or thin (\(m \times n\)) matrix.
LQDecomposition::execute(tm, pA, pQ, scratch): LQ of a wide matrix (\(m \leq n\)), computed by applying QR to the transpose.
SquareSVD::execute(tm, pA, pU, pV, sings, scratch, iscratch): SVD of a matrix with \(n_{rows} \geq n_{cols}\), computed by Householder bidiagonalization followed by implicit-shift QR iteration. The singular values are returned non-negative but are not sorted.
SymmetricEVD::execute(tm, pA, pQ, eigs, scratch, iscratch): eigen decomposition of a symmetric matrix, computed by Householder tridiagonalization followed by implicit QR. The eigenvalues are not sorted.

SquareSVD and SymmetricEVD return the number of QR iterations they performed. QRDecomposition and LQDecomposition return 0.

Workspace

Besides the matrices themselves, each routine needs a double workspace and, for SquareSVD and SymmetricEVD, a std::size_t workspace. Each class reports the required lengths through double_scratch_size(...) and sizet_scratch_size(...). total_shmem_scratch_size(...) returns the number of bytes to request from par_for_outer for both workspaces.

Matrix types

The routines are templated on the matrix type. Any type works if it provides operator()(int r, int c) and has GetNrows/GetNcols overloads that can be found by argument-dependent lookup. Rank-2 Kokkos views work directly. matrix_wrapper_t<T> wraps a row-major pointer, for example into scratch memory, and matrix_transpose_wrapper_t and the permuted-row/column wrappers in matrix_utils.hpp provide views without copies.

Ways to call the routines

The first argument of execute, tm, is the execution handle. It selects one of three ways to use the library.

On the host

Host-only overloads omit tm and the workspace arguments and allocate their own workspace. They are convenient for setup code and testing:

using namespace parthenon::batched_linear_algebra;
QRDecomposition::execute(&A, &Q); // host only, allocates scratch
QRDecomposition::execute(&A);     // R only

In a flat kernel (one thread per matrix)

Passing serial_tm_t() runs the whole decomposition on the calling thread, so it can be called from an ordinary par_for with one matrix per iteration. This suits large batches of very small matrices, where one thread per matrix already exposes enough parallelism. There is no team scratch in a flat kernel, so the matrices and workspace live either in per-thread local arrays, when the sizes are known at compile time, or in slices of preallocated device arrays:

using namespace parthenon::batched_linear_algebra;
constexpr int m = 6, n = 3;

parthenon::par_for(
    "BatchedQRFlat", 0, nbatch - 1, KOKKOS_LAMBDA(const int b) {
      constexpr std::size_t kWorkSize = QRDecomposition::double_scratch_size(m, n);
      double a_data[m * n], q_data[m * n], work[kWorkSize];
      matrix_wrapper_t<double> A(a_data, m, n);
      matrix_wrapper_t<double> Q(q_data, m, n);

      // Fill A for this batch entry, e.g. from cell data
      for (int r = 0; r < m; ++r) {
        for (int c = 0; c < n; ++c) {
          A(r, c) = data(b, r, c);
        }
      }

      QRDecomposition::execute(serial_tm_t(), &A, &Q, work);

      // A now holds R and Q holds the thin Q factor; use them here
    });

In a hierarchical kernel (one team per matrix)

Passing a team_mbr_t spreads each decomposition over the threads of a team, which suits larger matrices or smaller batches. All threads of the team must call execute together. The matrices and workspace would usually be allocated in team scratch. execute does not end with a team barrier, so add one before reading the results:

using namespace parthenon::batched_linear_algebra;
using parthenon::ScratchPad1D;
using parthenon::ScratchPad2D;
const int m = 12, n = 6;
const std::size_t nscratch = QRDecomposition::double_scratch_size(m, n);
// Workspace plus scratch for A (m x n) and a thin Q (m x n)
const std::size_t scratch_bytes = QRDecomposition::total_shmem_scratch_size(m, n) +
                                  2 * ScratchPad2D<double>::shmem_size(m, n);
constexpr int scratch_level = 0;

parthenon::par_for_outer(
    "BatchedQR", scratch_bytes, scratch_level, 0, nbatch - 1,
    KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b) {
      ScratchPad1D<double> work(member.team_scratch(scratch_level), nscratch);
      ScratchPad2D<double> A(member.team_scratch(scratch_level), m, n);
      ScratchPad2D<double> Q(member.team_scratch(scratch_level), m, n);

      // Fill A for this batch entry, e.g. from cell data
      parthenon::par_for_inner(member, 0, m - 1, 0, n - 1,
                               [&](const int r, const int c) { A(r, c) = data(b, r, c); });
      member.team_barrier();

      QRDecomposition::execute(member, &A, &Q, work.data());
      member.team_barrier();

      // A now holds R and Q holds the thin Q factor; use them here
    });