Skip to content

Conversation

@yasahi-hpc
Copy link
Contributor

This PR aims at improving the implementation and testing of serial/team/teamvector axpy.

  • Moving implementation details into Impl namespace.
  • If alpha is a View, we assume X and Y are 2D Views including the batch dimension. If alpha is a scalar, we assume X and Y are 1D Views without batch dimension. This aligns with other batched kernels.
Kokkos::parallel_for(
        "axpy", policy, KOKKOS_LAMBDA(int ib) {
          auto alpha = m_alpha(k);
          auto sub_x = Kokkos::subview(m_X, k, Kokkos::ALL);
          auto sub_y = Kokkos::subview(m_Y, k, Kokkos::ALL);
          // y = y + alpha * x
          KokkosBatched::SerialAxpy::invoke(alpha, sub_x, sub_y);
        });
  • Adding unit-tests for the scalar cases

Signed-off-by: Yuuichi Asahi <[email protected]>
@yasahi-hpc yasahi-hpc self-assigned this Oct 27, 2025
@yasahi-hpc yasahi-hpc added enhancement Cleanup Code maintenance that isn't a bugfix or new feature labels Oct 27, 2025
@yasahi-hpc
Copy link
Contributor Author

Should I need to update docs in this PR?
I am planning to first add an example and add docs.
Would that be fine with you?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Cleanup Code maintenance that isn't a bugfix or new feature enhancement

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants