KokkosBatched::Pttrf

Defined in header: KokkosBatched_Pttrf.hpp

template <typename ArgAlgo>
struct SerialPttrf {
  template <typename DViewType, typename EViewType>
  KOKKOS_INLINE_FUNCTION
  static int
  invoke(const DViewType& d,
         const EViewType& e);
};

Computes a factorization of a tridiagonal matrix A.

1. For a symmetric positive definite tridiagonal matrix A, this computes the \(A = L \cdot D \cdot L^T\) factorization. This operation is equivalent to the LAPACK routine SPTTRF or DPTTRF for single or double precision.

2. For a complex Hermitian positive definite tridiagonal matrix A, this computes the \(A = L \cdot D \cdot L^H\) (or \(A = U^H \cdot D \cdot U\)) factorization. This operation is equivalent to the LAPACK routine CPTTRF or ZPTTRF for single or double precision.

Parameters

d:

Input view containing the diagonal elements of D

e:

Input view containing the subdiagonal (superdiagonal) elements of L (U)

Type Requirements

• ArgAlgo must be KokkosBatched::Algo::Pttrf::Unblocked for the unblocked algorithm

• DViewType must be a Kokkos View of rank 1 containing the diagonal elements (length n)

• EViewType must be a Kokkos View of rank 1 containing the superdiagonal or subdiagonal elements (length n-1)

Example

//@HEADER
// ************************************************************************
//
//                        Kokkos v. 4.0
//       Copyright (2022) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
// See https://kokkos.org/LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//@HEADER

#include <Kokkos_Core.hpp>
#include <Kokkos_Random.hpp>
#include <KokkosBatched_Pttrf.hpp>
#include <KokkosBatched_Pttrs.hpp>

using ExecutionSpace = Kokkos::DefaultExecutionSpace;

/// \brief Example of batched pttrf/pttrs
/// Solving A * x = b, where
///   A: [[4, 1, 0],
///       [1, 4, 1],
///       [0, 1, 4]]
///   b: [1, 1, 1]
///   x: [3/14, 1/7, 3/14]
/// In tridiagonal storage,
///   d: [4, 4, 4] (diagonal elements)
///   e: [1, 1] (sub/super-diagonal elements)
///
/// This corresponds to the following system of equations:
///        4 x0 +   x1        = 1
///          x0 + 4 x1 +   x2 = 1
///                 x1 + 4 x2 = 1
///
int main(int /*argc*/, char** /*argv*/) {
  Kokkos::initialize();
  {
    using View2DType = Kokkos::View<double**, ExecutionSpace>;
    const int Nb = 10, n = 3;

    // Matrix A in tridiagonal storage
    View2DType d("d", Nb, n), e("e", Nb, n - 1);

    // Solution
    View2DType x("x", Nb, n);

    // Initialize d, e, and x with deep_copy
    const double d0 = 4.0, e0 = 1.0, x0 = 1.0;
    Kokkos::deep_copy(d, d0);
    Kokkos::deep_copy(e, e0);
    Kokkos::deep_copy(x, x0);

    // solve A * x = b with pttrf and pttrs
    ExecutionSpace exec;
    using policy_type = Kokkos::RangePolicy<ExecutionSpace, Kokkos::IndexType<int>>;
    policy_type policy{exec, 0, Nb};
    Kokkos::parallel_for(
        "pttrf-pttrs", policy, KOKKOS_LAMBDA(int ib) {
          auto sub_d = Kokkos::subview(d, ib, Kokkos::ALL);
          auto sub_e = Kokkos::subview(e, ib, Kokkos::ALL);
          auto sub_x = Kokkos::subview(x, ib, Kokkos::ALL);

          // Factorize d and e by pttrf
          KokkosBatched::SerialPttrf<KokkosBatched::Algo::Pttrf::Unblocked>::invoke(sub_d, sub_e);

          // Solve A * x = b with pttrs
          KokkosBatched::SerialPttrs<KokkosBatched::Uplo::Upper, KokkosBatched::Algo::Pttrs::Unblocked>::invoke(
              sub_d, sub_e, sub_x);
        });

    // Confirm that the results are correct
    auto h_x     = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, x);
    bool correct = true;
    double eps   = 1.0e-12;
    for (int ib = 0; ib < Nb; ib++) {
      if (Kokkos::abs(h_x(ib, 0) - 3.0 / 14.0) > eps) correct = false;
      if (Kokkos::abs(h_x(ib, 1) - 1.0 / 7.0) > eps) correct = false;
      if (Kokkos::abs(h_x(ib, 2) - 3.0 / 14.0) > eps) correct = false;
    }

    if (correct) {
      std::cout << "pttrf/pttrs works correctly!" << std::endl;
    }
  }
  Kokkos::finalize();
}

output:

pttrf/pttrs works correctly!