KokkosBatched::Axpy

Defined in header: KokkosBatched_Axpy.hpp

struct SerialAxpy {
  template <typename XViewType, typename YViewType, typename alphaViewType>
  KOKKOS_INLINE_FUNCTION static int invoke(const alphaViewType &alpha, const XViewType &X, const YViewType &Y);
};

template <typename MemberType>
struct TeamAxpy {
  template <typename XViewType, typename YViewType, typename alphaViewType>
  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const alphaViewType &alpha, const XViewType &X,
                                           const YViewType &Y);
};

template <typename MemberType>
struct TeamVectorAxpy {
  template <typename XViewType, typename YViewType, typename alphaViewType>
  KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const alphaViewType &alpha, const XViewType &X,
                                           const YViewType &Y);
};

Add entries of \(X\) scaled by \(\alpha\) to entries of \(Y\)

\[\begin{align} Y &= Y + \alpha X \end{align}\]

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

2. For complex vectors \(X\) and \(Y\), this operation is equivalent to the BLAS routine CAXPY or ZAXPY for single or double precision.

Parameters

alpha:

A scalar scaling factor or a length \(m\) vector of scaling factors.

x:

On input, \(X\) is a length \(n\) vector or a \(m\) by \(n\) matrix.

y:

On input, \(Y\) is a length \(n\) vector or a \(m\) by \(n\) matrix. On output, \(Y\) is overwritten by the updated vector or matrix.

Type Requirements

• MemberType must be a Kokkos team member handle (only for TeamAxpy and TeamVectorAxpy).

• XViewType must be a Kokkos View of rank 1 or 2 containing a vector \(X\)

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

• alphaViewType must be either a Kokkos View of rank 1 containing a matrix \(\alpha\) or a built-in arithmetic type like float, double, Kokkos::complex<float>, or Kokkos::complex<double>

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 axpy). When the input views \(X\) and \(Y\) are of rank 2, the kernel performs a matrix operation where each row of the matrices is treated as a separate vector.

Example

// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project

#include <Kokkos_Core.hpp>
#include <Kokkos_Random.hpp>
#include <KokkosBatched_Axpy.hpp>

using ExecutionSpace = Kokkos::DefaultExecutionSpace;

/// \brief Example of batched axpy
/// Add entries of X scaled by coeffecient a to entries of Y: Y += a*X
///   x: [1, 2, 3, 4]
///   y: [5, 6, 7, 8]
///   a: 1.5
///
/// After, y = y + a * x, it will give
///   y: [6.5, 9.0, 11.5, 14.0]
///
int main(int /*argc*/, char** /*argv*/) {
  Kokkos::initialize();
  {
    using View1DType = Kokkos::View<double*, ExecutionSpace>;
    using View2DType = Kokkos::View<double**, ExecutionSpace>;
    const int Nb = 10, n = 4;

    // Scalar alpha
    View1DType alpha("alpha", Nb);

    // Vector x
    View2DType x("x", Nb, n), y("y", Nb, n);

    // Initialize alpha, x and y
    Kokkos::deep_copy(alpha, 1.5);
    auto h_x = Kokkos::create_mirror_view(x);
    auto h_y = Kokkos::create_mirror_view(y);
    for (int ib = 0; ib < Nb; ib++) {
      // Fill vector x and y
      for (int i = 0; i < n; i++) {
        h_x(ib, i) = i + 1;
        h_y(ib, i) = i + 5;
      }
    }
    Kokkos::deep_copy(x, h_x);
    Kokkos::deep_copy(y, h_y);

    // Compute y = y + a * x
    ExecutionSpace exec;
    using policy_type = Kokkos::RangePolicy<ExecutionSpace, Kokkos::IndexType<int>>;
    policy_type policy{exec, 0, Nb};
    Kokkos::parallel_for(
        "axpy", policy, KOKKOS_LAMBDA(int ib) {
          // y = y + a * x
          auto a     = alpha(ib);
          auto sub_x = Kokkos::subview(x, ib, Kokkos::ALL());
          auto sub_y = Kokkos::subview(y, ib, Kokkos::ALL());
          KokkosBatched::SerialAxpy::invoke(a, sub_x, sub_y);
        });

    // Confirm that the results are correct
    Kokkos::deep_copy(h_y, y);
    bool correct = true;
    double eps   = 1.0e-12;
    for (int ib = 0; ib < Nb; ib++) {
      if (Kokkos::abs(h_y(ib, 0) - 6.5) > eps) correct = false;
      if (Kokkos::abs(h_y(ib, 1) - 9.0) > eps) correct = false;
      if (Kokkos::abs(h_y(ib, 2) - 11.5) > eps) correct = false;
      if (Kokkos::abs(h_y(ib, 3) - 14.0) > eps) correct = false;
    }

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

output:

axpy works correctly!