team_ops.hpp

Go to the documentation of this file.

//
// This file is subject to the terms and conditions defined in
// file 'LICENSE', which is part of this source code package.
//
 
#pragma once
 
#include <Kokkos_Core.hpp>
 
#include "common/concepts.hpp"
#include "common/types.hpp"
 
namespace PACMAN {
 
namespace RbfPum {

template <typename TeamHandleType, KokkosViewRank<2> AViewType,
          KokkosViewRank<1> XType, KokkosViewRank<1> XOffsType,
          RBFFunction RbfFuncType>
KOKKOS_FORCEINLINE_FUNCTION void
TeamFillA(const TeamHandleType &teamHandle, const AViewType &A, const XType &X,
          const XOffsType &XOffs, const RbfFuncType &func) {
  const int_t N = A.extent_int(0);
  Kokkos::parallel_for(
      Kokkos::TeamThreadRange(teamHandle, N), [&](const int_t &i) {
        const auto x = X(XOffs(i));
        Kokkos::parallel_for(
            Kokkos::ThreadVectorRange(teamHandle, i, N), [&](const int &j) {
              const auto val = func.Eval(DistanceNoCheck(x, X(XOffs(j))));
              A(i, j) = val;
              A(j, i) = val;
            });
      });
}

template <typename TeamHandleType, KokkosViewRank<2> AViewType,
          KokkosViewRank<1> XType, KokkosViewRank<1> XOffsType,
          KokkosViewRank<1> YType, KokkosViewRank<1> YOffsType,
          RBFFunction RbfFuncType>
KOKKOS_FORCEINLINE_FUNCTION void
TeamFillE(const TeamHandleType &teamHandle, AViewType &A, const XType &X,
          const XOffsType &XOffs, const YType &Y, const YOffsType &YOffs,
          const RbfFuncType &func) {
  const int_t N = A.extent_int(0);
  const int_t M = A.extent_int(1);
 
  Kokkos::parallel_for(
      Kokkos::TeamThreadRange(teamHandle, N), [&](const int_t &i) {
        const auto y = Y(YOffs(i));
        Kokkos::parallel_for(
            Kokkos::ThreadVectorRange(teamHandle, M), [&](const int_t &j) {
              A(i, j) = func.Eval(DistanceNoCheck(y, X(XOffs(j))));
            });
      });
}

template <typename TeamHandleType, KokkosViewRank<1> BViewType,
          KokkosViewRank<1> XType, KokkosViewRank<1> XOffsType>
KOKKOS_FORCEINLINE_FUNCTION void TeamFillB(const TeamHandleType &teamHandle,
                                           BViewType &B, const XType &X,
                                           const XOffsType &XOffs) {
  const int_t N = B.extent_int(0);
  Kokkos::parallel_for(Kokkos::TeamVectorRange(teamHandle, N),
                       [&](const int_t &i) { B(i) = X(XOffs(i)); });
}
 
} // namespace RbfPum
 
} // namespace PACMAN