linear_interpolator_1d.hh

Go to the documentation of this file.

#pragma once
 
// DiFfRG
#include <DiFfRG/common/kokkos.hh>
#include <DiFfRG/common/math.hh>
#include <DiFfRG/physics/interpolation/interpolation_stencil.hh>
 
// std
#include <cstring>
 
namespace DiFfRG
{
  template <typename NT, typename Coordinates> class LinearInterpolator1D
  {
    static_assert(Coordinates::dim == 1, "LinearInterpolator1D requires 1D coordinates");
 
    static constexpr bool periodic = is_periodic_coordinate_v<Coordinates>;
 
    using ViewType = Kokkos::View<NT *, GPU_memory, Kokkos::MemoryTraits<Kokkos::RandomAccess>>;
    using HostViewType = typename ViewType::host_mirror_type;
 
    // Without a distinct device space create_mirror_view aliases the source: the two members are
    // then one allocation, and the upload must not be compiled at all so a host-only build pays no
    // fence.
    static constexpr bool has_separate_device =
        !std::is_same_v<typename ViewType::memory_space, typename HostViewType::memory_space>;
 
  public:
    using ctype = typename Coordinates::ctype;
    using value_type = NT;
    static constexpr size_t dim = 1;
 
    LinearInterpolator1D(const Coordinates &coordinates)
        : coordinates(coordinates), size(coordinates.size()),
          device_data("LinearInterpolator1D_data", coordinates.size()),
          host_data(Kokkos::create_mirror_view(device_data))
    {
    }
 
    KOKKOS_DEFAULTED_FUNCTION LinearInterpolator1D(const LinearInterpolator1D &) = default;
 
    template <typename NT2> void update(const NT2 *in_data)
    {
      if constexpr (std::is_same_v<NT, NT2> && std::is_trivially_copyable_v<NT>)
        std::memcpy(host_data.data(), in_data, size * sizeof(NT));
      else
        for (size_t i = 0; i < size; ++i)
          host_data(i) = in_data[i];
 
      if constexpr (has_separate_device) {
        typename ViewType::execution_space exec;
        Kokkos::deep_copy(exec, device_data, host_data);
        exec.fence();
      }
    }
 
    NT operator[](size_t i) const { return host_data(i); }
 
    typename Coordinates::ctype KOKKOS_FUNCTION index(const typename Coordinates::ctype x) const
    {
      return coordinates.backward(x);
    }
 
    NT KOKKOS_FUNCTION at(const typename Coordinates::ctype idx) const
    {
      // Resolve the stencil: clamped [i, i+1] for a bounded axis, wrapping [i, (i+1) % size] for a periodic one
      const auto stencil = make_interpolation_stencil<periodic>(idx, size);
      const auto t = stencil.t;
      const auto lower = value(stencil.lower);
      const auto upper = value(stencil.upper);
      if constexpr (std::is_arithmetic_v<NT>)
        return Kokkos::fma(t, upper, Kokkos::fma(-t, lower, lower));
      else
        return t * upper + (1 - t) * lower;
    }
 
    NT KOKKOS_FUNCTION operator()(const typename Coordinates::ctype x) const { return at(index(x)); }
 
    const Coordinates &get_coordinates() const { return coordinates; }
 
    const NT *data() const { return host_data.data(); }
 
  private:
    KOKKOS_FORCEINLINE_FUNCTION NT value(const size_t i) const
    {
      KOKKOS_IF_ON_DEVICE((return device_data(i);))
      KOKKOS_IF_ON_HOST((return host_data(i);))
    }
 
    const Coordinates coordinates;
    const size_t size;
 
    ViewType device_data;
    HostViewType host_data;
  };
} // namespace DiFfRG