KokkosBatched::Swap

Defined in header: KokkosBatched_Swap.hpp

struct SerialSwap {
  template <typename XViewType, typename YViewType>
  KOKKOS_INLINE_FUNCTION static int invoke(const XViewType &x, const YViewType &y);
};

template <typename MemberType>
struct TeamSwap {
  template <typename XViewType, typename YViewType>
  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const XViewType &x, const YViewType &y);
};

template <typename MemberType>
struct TeamVectorSwap {
  template <typename XViewType, typename YViewType>
  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const XViewType &x, const YViewType &y);
};

Swaps the elements of two vectors or matrices \(x\) and \(y\).

  • For a real vector or matrix \(x\), this operation is equivalent to the BLAS routine SSWAP or DSWAP for single or double precision.

  • For a complex vector or matrix \(x\), this operation is equivalent to the BLAS routine CSWAP or ZSWAP for single or double precision.

Parameters

x:

\(x\) is a length \(n\) vector or a \(m\) by \(n\) matrix.

y:

\(y\) is a length \(n\) vector or a \(m\) by \(n\) matrix.

Type Requirements

  • MemberType must be a Kokkos team member handle (only for TeamSwap and TeamVectorSwap)

  • XViewType must be a Kokkos View of rank 1 or 2 containing a vector or matrix \(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 or 2 containing a vector or matrix \(y\) that satisfies std::is_same_v<typename YViewType::value_type, typename YViewType::non_const_value_type>

  • The value type of XViewType and YViewType must be swappable.

Note

This kernel supports both vector and matrix operations. When the input views \(X\) and \(Y\) are of rank 1, the kernel performs a vector operation (BLAS swap). When the input views \(X\) and \(Y\) are of rank 2, the kernel swaps matrices \(X\) and \(Y\).

Example

 1#include <iostream>
 2#include <Kokkos_Core.hpp>
 3#include <KokkosBatched_Swap.hpp>
 4
 5using ExecutionSpace = Kokkos::DefaultExecutionSpace;
 6
 7/// \brief Example of batched swap
 8/// Swapping a batch of two vectors x and y
 9///
10/// Usage example:
11///   Before: x: [1, 3, 5], y: [2, 4, 6]
12///   After:  x: [2, 4, 6], y: [1, 3, 5]
13///
14int main(int /*argc*/, char** /*argv*/) {
15  Kokkos::initialize();
16  {
17    using View2DType = Kokkos::View<double**, ExecutionSpace>;
18    const int Nb = 10, n = 3;
19
20    // Vectors x and y
21    View2DType x("x", Nb, n), y("y", Nb, n);
22
23    // Initialize x and y
24    auto h_x = Kokkos::create_mirror_view(x);
25    auto h_y = Kokkos::create_mirror_view(y);
26    for (int ib = 0; ib < Nb; ib++) {
27      h_x(ib, 0) = 1;
28      h_x(ib, 1) = 3;
29      h_x(ib, 2) = 5;
30      h_y(ib, 0) = 2;
31      h_y(ib, 1) = 4;
32      h_y(ib, 2) = 6;
33    }
34    Kokkos::deep_copy(x, h_x);
35    Kokkos::deep_copy(y, h_y);
36
37    // Swap x and y
38    ExecutionSpace exec;
39    using policy_type = Kokkos::RangePolicy<ExecutionSpace, Kokkos::IndexType<int>>;
40    policy_type policy{exec, 0, Nb};
41    Kokkos::parallel_for(
42        "swap", policy, KOKKOS_LAMBDA(int ib) {
43          auto sub_x = Kokkos::subview(x, ib, Kokkos::ALL());
44          auto sub_y = Kokkos::subview(y, ib, Kokkos::ALL());
45          KokkosBatched::SerialSwap::invoke(sub_x, sub_y);
46        });
47
48    // Confirm that the results are correct
49    Kokkos::deep_copy(h_x, x);
50    Kokkos::deep_copy(h_y, y);
51    bool correct = true;
52    double eps   = 1.0e-12;
53    for (int ib = 0; ib < Nb; ib++) {
54      if (Kokkos::abs(h_x(ib, 0) - 2) > eps) correct = false;
55      if (Kokkos::abs(h_x(ib, 1) - 4) > eps) correct = false;
56      if (Kokkos::abs(h_x(ib, 2) - 6) > eps) correct = false;
57      if (Kokkos::abs(h_y(ib, 0) - 1) > eps) correct = false;
58      if (Kokkos::abs(h_y(ib, 1) - 3) > eps) correct = false;
59      if (Kokkos::abs(h_y(ib, 2) - 5) > eps) correct = false;
60    }
61
62    if (correct) {
63      std::cout << "swap works correctly!" << std::endl;
64    }
65  }
66  Kokkos::finalize();
67}

output:

swap works correctly!