KokkosBatched::Iamax

Defined in header: KokkosBatched_Iamax.hpp

struct SerialIamax {
  template <typename XViewType>
  KOKKOS_INLINE_FUNCTION static typename XViewType::size_type invoke(const XViewType &x);
};

template <typename MemberType>
struct TeamIamax {
  template <typename XViewType>
  KOKKOS_INLINE_FUNCTION static typename XViewType::size_type invoke(const MemberType &member, const XViewType &x);
};

template <typename MemberType>
struct TeamVectorIamax {
  template <typename XViewType>
  KOKKOS_INLINE_FUNCTION static typename XViewType::size_type invoke(const MemberType &member, const XViewType &x);
};

Finds the index of the first element of \(x\) having maximum absolute value. As well as Blas, this returns 0 for an empty vector.

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

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

Parameters

x:

\(x\) is a length n vector.

Type Requirements

  • MemberType must be a Kokkos team member handle (only for TeamIamax and TeamVectorIamax)

  • XViewType must be a Kokkos View of rank 1 containing a vector \(x\)

Example

 1#include <Kokkos_Core.hpp>
 2#include <Kokkos_Random.hpp>
 3#include <KokkosBatched_Iamax.hpp>
 4
 5using ExecutionSpace = Kokkos::DefaultExecutionSpace;
 6
 7/// \brief Example of batched iamax
 8/// Finds the index of the first element having maximum absolute value.
 9///   X0: [1, 2, 0] -> 1
10///   X1: [-5, 4, 3] -> 0
11///   X2: [0, -1, -1] -> 1
12///
13int main(int /*argc*/, char** /*argv*/) {
14  Kokkos::initialize();
15  {
16    using View2DType = Kokkos::View<double**, ExecutionSpace>;
17    using Idx1DType  = Kokkos::View<int*, ExecutionSpace>;
18    const int Nb = 10, n = 3;
19
20    // Batched vectors
21    View2DType x0("x0", Nb, n), x1("x1", Nb, n), x2("x2", Nb, n);
22
23    // Max indices
24    Idx1DType iamax0("iamax0", Nb), iamax1("iamax1", Nb), iamax2("iamax2", Nb);
25
26    // Initialize x0, x1, x2
27    auto h_x0 = Kokkos::create_mirror_view(x0);
28    auto h_x1 = Kokkos::create_mirror_view(x1);
29    auto h_x2 = Kokkos::create_mirror_view(x2);
30
31    for (int ib = 0; ib < Nb; ib++) {
32      h_x0(ib, 0) = 1.0;
33      h_x0(ib, 1) = 2.0;
34      h_x0(ib, 2) = 0.0;
35
36      h_x1(ib, 0) = -5.0;
37      h_x1(ib, 1) = 4.0;
38      h_x1(ib, 2) = 3.0;
39
40      h_x2(ib, 0) = 0.0;
41      h_x2(ib, 1) = -1.0;
42      h_x2(ib, 2) = -1.0;
43    }
44    Kokkos::deep_copy(x0, h_x0);
45    Kokkos::deep_copy(x1, h_x1);
46    Kokkos::deep_copy(x2, h_x2);
47
48    // Find max indices
49    ExecutionSpace exec;
50    using policy_type = Kokkos::RangePolicy<ExecutionSpace, Kokkos::IndexType<int>>;
51    policy_type policy{exec, 0, Nb};
52    Kokkos::parallel_for(
53        "iamax", policy, KOKKOS_LAMBDA(int ib) {
54          auto sub_x0 = Kokkos::subview(x0, ib, Kokkos::ALL);
55          auto sub_x1 = Kokkos::subview(x1, ib, Kokkos::ALL);
56          auto sub_x2 = Kokkos::subview(x2, ib, Kokkos::ALL);
57
58          // Find max indices
59          iamax0(ib) = KokkosBatched::SerialIamax::invoke(sub_x0);
60          iamax1(ib) = KokkosBatched::SerialIamax::invoke(sub_x1);
61          iamax2(ib) = KokkosBatched::SerialIamax::invoke(sub_x2);
62        });
63
64    // Confirm that the results are correct
65    auto h_iamax0 = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, iamax0);
66    auto h_iamax1 = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, iamax1);
67    auto h_iamax2 = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, iamax2);
68    bool correct  = true;
69    for (int ib = 0; ib < Nb; ib++) {
70      if (h_iamax0(ib) != 1) correct = false;
71      if (h_iamax1(ib) != 0) correct = false;
72      if (h_iamax2(ib) != 1) correct = false;
73    }
74
75    if (correct) {
76      std::cout << "iamax works correctly!" << std::endl;
77    }
78  }
79  Kokkos::finalize();
80}

output:

iamax works correctly!