KokkosBatched::Rot

Defined in header: KokkosBatched_Rot.hpp

template <bool Conj = false>
struct SerialRot {
  template <typename XViewType, typename YViewType, typename CType, typename SType>
  KOKKOS_INLINE_FUNCTION static int invoke(const XViewType &x, const YViewType &y, const CType c, const SType s);
};

template <typename MemberType, bool Conj = false>
struct TeamAxpy {
  template <typename XViewType, typename YViewType, typename CType, typename SType>
  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const XViewType &x, const YViewType &y,
                                           const CType c, const SType s);
};

template <typename MemberType, bool Conj = false>
struct TeamVectorRot {
  template <typename XViewType, typename YViewType, typename CType, typename SType>
  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const XViewType &x, const YViewType &y,
                                           const CType c, const SType s);
};

Applies a plane rotation to vectors \(x\) and \(y\):

\[\begin{split}\begin{align} x &= c * x + s * y \\ y &= -s * x + c * y \: \text{(if Conj is false)} \\ y &= -s * conj(x) + c * y \: \text{(if Conj is true)} \end{align}\end{split}\]

  1. For real vectors \(X\) and \(Y\), this operation is equivalent to the BLAS routine SROT or DROT for single or double precision.

  2. For complex vectors \(X\) and \(Y\), this operation is equivalent to the BLAS routine CSROT or ZDROT for single or double precision if Conj is false. If Conj is true, this operation is equivalent to the BLAS routine CROT or ZROT for single or double precision.

Parameters

x:

On input, \(x\) is a length \(n\) vector. On output, \(x\) is overwritten by the rotated vector.

y:

On input, \(y\) is a length \(n\) vector. On output, \(y\) is overwritten by the rotated vector.

c:

A scalar of cosine of the rotation (real scalar)

s:

A scalar of sine of the rotation (real or complex scalar)

Type Requirements

  • Conj must be a boolean template parameter that indicates whether the rotation is a conjugate rotation or not.

  • MemberType must be a Kokkos team member handle (only for TeamRot and TeamVectorRot).

  • XViewType must be a Kokkos View of rank 1 containing a vector \(X\) that satisfies std::is_same_v<typename XViewType::value_type, typename XViewType::non_const_value_type>

  • YViewType must be a Kokkos View of rank 1 containing a vector \(Y\) that satisfies std::is_same_v<typename YViewType::value_type, typename YViewType::non_const_value_type>

  • CType must be a built-in real type like float, or double

  • SType must be a built-in arithmetic type like float, double, Kokkos::complex<float>, or Kokkos::complex<double>

Example

 1#include <Kokkos_Core.hpp>
 2#include <KokkosBatched_Rot.hpp>
 3#include <KokkosBatched_Rotg.hpp>
 4
 5using ExecutionSpace = Kokkos::DefaultExecutionSpace;
 6
 7/// \brief Example of batched rot and rotg
 8/// Compute c ans s such that
 9/// [[c,  s],  * [[a],  = [[r],
10///  [-s, c]]     [b]]     [0]]
11/// Then apply the rotation to vectors X and Y:
12/// [[X'],  = [[c,  s],  * [[X],
13///  [Y']]     [-s, c]]     [Y]]
14///
15/// Usage example:
16/// a = 3, b = 4
17/// X = [3, 3, 3], Y = [4, 4, 4]
18/// X' = [5, 5, 5], Y' = [0, 0, 0]
19///
20int main(int /*argc*/, char** /*argv*/) {
21  Kokkos::initialize();
22  {
23    using View1DType = Kokkos::View<double*, ExecutionSpace>;
24    using View2DType = Kokkos::View<double**, ExecutionSpace>;
25    const int Nb = 10, n = 3;
26
27    // Values A and B we want to rotate
28    View1DType A("A", Nb), B("B", Nb), C("C", Nb), S("S", Nb);
29
30    // Vector X and Y of the rot/rotg operation
31    View2DType X("X", Nb, n), Y("Y", Nb, n);
32
33    // Initialize A, B, X and Y
34    const double a = 3.0, b = 4.0;
35    Kokkos::deep_copy(A, a);
36    Kokkos::deep_copy(B, b);
37    Kokkos::deep_copy(X, a);
38    Kokkos::deep_copy(Y, b);
39
40    // Compute givens rotation coefficients and apply the rotation to X and Y
41    ExecutionSpace exec;
42    using policy_type = Kokkos::RangePolicy<ExecutionSpace, Kokkos::IndexType<int>>;
43    policy_type policy{exec, 0, Nb};
44    Kokkos::parallel_for(
45        "rot/rotg", policy, KOKKOS_LAMBDA(int ib) {
46          auto sub_A = Kokkos::subview(A, ib);
47          auto sub_B = Kokkos::subview(B, ib);
48          auto sub_C = Kokkos::subview(C, ib);
49          auto sub_S = Kokkos::subview(S, ib);
50          KokkosBatched::Rotg::invoke(sub_A, sub_B, sub_C, sub_S);
51
52          auto sub_X = Kokkos::subview(X, ib, Kokkos::ALL());
53          auto sub_Y = Kokkos::subview(Y, ib, Kokkos::ALL());
54          KokkosBatched::SerialRot<>::invoke(sub_X, sub_Y, sub_C(), sub_S());
55        });
56
57    // Confirm that the results are correct
58    auto h_X     = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, X);
59    auto h_Y     = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, Y);
60    bool correct = true;
61    double eps   = 1.0e-12;
62    for (int ib = 0; ib < Nb; ib++) {
63      for (int j = 0; j < n; j++) {
64        if (Kokkos::abs(h_X(ib, j) - 5) > eps) correct = false;
65        if (Kokkos::abs(h_Y(ib, j)) > eps) correct = false;
66      }
67    }
68
69    if (correct) {
70      std::cout << "rot/rotg works correctly!" << std::endl;
71    }
72  }
73  Kokkos::finalize();
74}

output:

rot/rotg works correctly!