KokkosBatched::Rotmg

Defined in header: KokkosBatched_Rotmg.hpp

struct Rotmg {
  template <class DXViewType, class YViewType, class PViewType>
  KOKKOS_INLINE_FUNCTION static int invoke(const DXViewType &d1, const DXViewType &d2, const DXViewType &x1,
                                           const YViewType &y1, const PViewType &param);
};

Constructs the modified Givens transformation

\[\begin{split}H = \begin{bmatrix} h_{11} & h_{12} \\ h_{21} & h_{22} \end{bmatrix}\end{split}\]

that zeros out the second component of a 2D vector \((\sqrt{d1} \times x1, \sqrt{d2} \times y1)^T\). The transformation can be applied to vectors \(x\) and \(y\) by calling the Rotm routine described in KokkosBatched::Rotm. This operation is equivalent to the BLAS routine SROTMG or DROTMG for single or double precision. A \(flag\) and matrix H are stored in a 5-element vector \(param\) as follows:

\[\begin{split}\begin{align} param(0) &= flag \\ param(1) &= h_{11} \\ param(2) &= h_{21} \\ param(3) &= h_{12} \\ param(4) &= h_{22} \end{align}\end{split}\]

The value of \(flag\) can be -1, 0, 1, or -2, which corresponds to the following forms of H:

\(flag == -1\)

\(flag == 0\)

\(flag == 1\)

\(flag == -2\)

\(\begin{bmatrix} h_{11} & h_{12} \\ h_{21} & h_{22} \end{bmatrix}\)

\(\begin{bmatrix} 1.0 & h_{12} \\ h_{21} & 1.0 \end{bmatrix}\)

\(\begin{bmatrix} h_{11} & 1.0 \\ -1.0 & h_{22} \end{bmatrix}\)

\(\begin{bmatrix} 1.0 & 0.0 \\ 0.0 & 1.0 \end{bmatrix}\)

Parameters

d1, d2:

0-D view. On input, the squares of the initial scaling factors for \(x\) and \(y\). On output, they contains squares of the scaling factors to be applied to \(x\) and \(y\), respectively.

x1, y1:

0-D view. On input, the components of the 2D vector to rotate. On output, \(x1\) contains the first component of the rotated vector, while \(y1\) is unchanged on output.

param:

0-D view. On output, A length 5 vector, a rank 1 view, containing the parameters of the modified Givens rotation. The first element of \(param\) is a flag that indicates the form of the rotation matrix H, and the remaining four elements are the entries of the rotation matrix H.

Type Requirements

  • DXViewType must be a Kokkos View of rank 0 that satisfies:

    • std::is_same_v<typename DXViewType::non_const_value_type, typename DXViewType::value_type> == true

    • !KokkosKernels::ArithTraits<typename DXViewType::value_type>::is_complex

  • YViewType must be a Kokkos View of rank 0 that satisfies:

    • !KokkosKernels::ArithTraits<typename YViewType::value_type>::is_complex

  • PViewType must be a Kokkos View of rank 1 that satisfies:

    • std::is_same_v<typename PViewType::non_const_value_type, typename PViewType::value_type> == true

    • !KokkosKernels::ArithTraits<typename PViewType::value_type>::is_complex

Example

 1#include <iostream>
 2#include <Kokkos_Core.hpp>
 3#include <KokkosBatched_Rotm.hpp>
 4#include <KokkosBatched_Rotmg.hpp>
 5
 6using ExecutionSpace = Kokkos::DefaultExecutionSpace;
 7
 8/// \brief Example of batched rotm and rotmg
 9/// rotmg constructs the modified Givens rotation matrix H which eliminates the second component of a 2-vector [d1, d2]
10/// and applies the rotation to a 2-vector [x1, y1].
11/// Then apply the rotation to vectors X and Y:
12/// [[X'],  = H  * [[X],
13///  [Y']]          [Y]]
14///
15/// Usage example:
16/// d1 = 1.0, d2 = 1.0, x1 = 3.0, y1 = 4.0
17/// X = [3, 3, 3], Y = [4, 4, 4]
18/// X' = [6.25, 6.25, 6.25], 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 d1 and d2 we want to rotate
28    View1DType d1("d1", Nb), d2("d2", Nb), x1("x1", Nb), y1("y1", Nb);
29
30    // Vector X and Y of the rot/rotg operation
31    View2DType param("param", Nb, 5);
32    View2DType X("X", Nb, n), Y("Y", Nb, n);
33
34    // Initialize d1, d2, X and Y
35    const double a = 3.0, b = 4.0;
36    Kokkos::deep_copy(d1, 1.0);
37    Kokkos::deep_copy(d2, 1.0);
38    Kokkos::deep_copy(x1, a);
39    Kokkos::deep_copy(y1, b);
40    Kokkos::deep_copy(X, a);
41    Kokkos::deep_copy(Y, b);
42
43    // Compute a modified Givens rotation coefficients and apply the rotation to X and Y
44    ExecutionSpace exec;
45    using policy_type = Kokkos::RangePolicy<ExecutionSpace, Kokkos::IndexType<int>>;
46    policy_type policy{exec, 0, Nb};
47    Kokkos::parallel_for(
48        "rotm/rotmg", policy, KOKKOS_LAMBDA(int ib) {
49          auto sub_d1    = Kokkos::subview(d1, ib);
50          auto sub_d2    = Kokkos::subview(d2, ib);
51          auto sub_x1    = Kokkos::subview(x1, ib);
52          auto sub_y1    = Kokkos::subview(y1, ib);
53          auto sub_param = Kokkos::subview(param, ib, Kokkos::ALL());
54          KokkosBatched::Rotmg::invoke(sub_d1, sub_d2, sub_x1, sub_y1, sub_param);
55
56          auto sub_X = Kokkos::subview(X, ib, Kokkos::ALL());
57          auto sub_Y = Kokkos::subview(Y, ib, Kokkos::ALL());
58          KokkosBatched::SerialRotm::invoke(sub_X, sub_Y, sub_param);
59        });
60
61    // Confirm that the results are correct
62    auto h_X     = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, X);
63    auto h_Y     = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, Y);
64    bool correct = true;
65    double eps   = 1.0e-12;
66    for (int ib = 0; ib < Nb; ib++) {
67      for (int j = 0; j < n; j++) {
68        if (Kokkos::abs(h_X(ib, j) - 6.25) > eps) correct = false;
69        if (Kokkos::abs(h_Y(ib, j)) > eps) correct = false;
70      }
71    }
72
73    if (correct) {
74      std::cout << "rotm/rotmg works correctly!" << std::endl;
75    }
76  }
77  Kokkos::finalize();
78}

output:

rotm/rotmg works correctly!