AMReX_ParticleTransformation.H

Go to the documentation of this file.

#ifndef AMREX_PARTICLETRANSFORMATION_H_
#define AMREX_PARTICLETRANSFORMATION_H_
#include <AMReX_Config.H>
 
#include <AMReX_IntVect.H>
#include <AMReX_Box.H>
#include <AMReX_Gpu.H>
#include <AMReX_Print.H>
#include <AMReX_ParticleTile.H>
#include <AMReX_ParticleUtil.H>
 
namespace amrex
{
 
template <typename T_ParticleType, int NAR, int NAI>
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void copyParticle (const      ParticleTileData<T_ParticleType, NAR, NAI>& dst,
                   const ConstParticleTileData<T_ParticleType, NAR, NAI>& src,
                   int src_i, int dst_i) noexcept
{
    AMREX_ASSERT(dst.m_num_runtime_real == src.m_num_runtime_real);
    AMREX_ASSERT(dst.m_num_runtime_int  == src.m_num_runtime_int );
 
    if constexpr(!T_ParticleType::is_soa_particle) {
        dst.m_aos[dst_i] = src.m_aos[src_i];
    } else {
        dst.m_idcpu[dst_i] = src.m_idcpu[src_i];
    }
    if constexpr(NAR > 0) {
        for (int j = 0; j < NAR; ++j) {
            dst.m_rdata[j][dst_i] = src.m_rdata[j][src_i];
        }
    }
    for (int j = 0; j < dst.m_num_runtime_real; ++j) {
        dst.m_runtime_rdata[j][dst_i] = src.m_runtime_rdata[j][src_i];
    }
    if constexpr(NAI > 0) {
        for (int j = 0; j < NAI; ++j) {
            dst.m_idata[j][dst_i] = src.m_idata[j][src_i];
        }
    }
    for (int j = 0; j < dst.m_num_runtime_int; ++j) {
        dst.m_runtime_idata[j][dst_i] = src.m_runtime_idata[j][src_i];
    }
}
 
template <typename T_ParticleType, int NAR, int NAI>
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void copyParticle (const ParticleTileData<T_ParticleType, NAR, NAI>& dst,
                   const ParticleTileData<T_ParticleType, NAR, NAI>& src,
                   int src_i, int dst_i) noexcept
{
    AMREX_ASSERT(dst.m_num_runtime_real == src.m_num_runtime_real);
    AMREX_ASSERT(dst.m_num_runtime_int  == src.m_num_runtime_int );
 
    if constexpr(T_ParticleType::is_soa_particle) {
        dst.m_idcpu[dst_i] = src.m_idcpu[src_i];
    } else {
        dst.m_aos[dst_i] = src.m_aos[src_i];
    }
    for (int j = 0; j < NAR; ++j) {
        dst.m_rdata[j][dst_i] = src.m_rdata[j][src_i];
    }
    for (int j = 0; j < dst.m_num_runtime_real; ++j) {
        dst.m_runtime_rdata[j][dst_i] = src.m_runtime_rdata[j][src_i];
    }
    for (int j = 0; j < NAI; ++j) {
        dst.m_idata[j][dst_i] = src.m_idata[j][src_i];
    }
    for (int j = 0; j < dst.m_num_runtime_int; ++j) {
        dst.m_runtime_idata[j][dst_i] = src.m_runtime_idata[j][src_i];
    }
}
 
template <typename T_ParticleType, int NAR, int NAI>
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void swapParticle (const ParticleTileData<T_ParticleType, NAR, NAI>& dst,
                   const ParticleTileData<T_ParticleType, NAR, NAI>& src,
                   int src_i, int dst_i) noexcept
{
    AMREX_ASSERT(dst.m_num_runtime_real == src.m_num_runtime_real);
    AMREX_ASSERT(dst.m_num_runtime_int  == src.m_num_runtime_int );
 
    if constexpr(T_ParticleType::is_soa_particle) {
        amrex::Swap(src.m_idcpu[src_i], dst.m_idcpu[dst_i]);
    } else {
        amrex::Swap(src.m_aos[src_i], dst.m_aos[dst_i]);
    }
    for (int j = 0; j < NAR; ++j) {
        amrex::Swap(dst.m_rdata[j][dst_i], src.m_rdata[j][src_i]);
    }
    for (int j = 0; j < dst.m_num_runtime_real; ++j) {
        amrex::Swap(dst.m_runtime_rdata[j][dst_i], src.m_runtime_rdata[j][src_i]);
    }
    for (int j = 0; j < NAI; ++j) {
        amrex::Swap(dst.m_idata[j][dst_i], src.m_idata[j][src_i]);
    }
    for (int j = 0; j < dst.m_num_runtime_int; ++j) {
        amrex::Swap(dst.m_runtime_idata[j][dst_i], src.m_runtime_idata[j][src_i]);
    }
}
 
template <typename DstTile, typename SrcTile>
void copyParticles (DstTile& dst, const SrcTile& src) noexcept
{
    auto np = src.numParticles();
    copyParticles(dst, src, 0, 0, np);
}
 
template <typename DstTile, typename SrcTile, typename Index, typename N,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
void copyParticles (DstTile& dst, const SrcTile& src,
                    Index src_start, Index dst_start, N n) noexcept
{
    const auto src_data = src.getConstParticleTileData();
          auto dst_data = dst.getParticleTileData();
 
    AMREX_HOST_DEVICE_FOR_1D( n, i,
    {
        copyParticle(dst_data, src_data, src_start+i, dst_start+i);
    });
 
    Gpu::streamSynchronize();
}
 
template <typename DstTile, typename SrcTile, typename F>
void transformParticles (DstTile& dst, const SrcTile& src, F&& f) noexcept
{
    auto np = src.numParticles();
    transformParticles(dst, src, 0, 0, np, std::forward<F>(f));
}
 
template <typename DstTile, typename SrcTile, typename Index, typename N, typename F,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
void transformParticles (DstTile& dst, const SrcTile& src,
                         Index src_start, Index dst_start, N n, F const& f) noexcept
{
    const auto src_data = src.getConstParticleTileData();
          auto dst_data = dst.getParticleTileData();
 
    AMREX_HOST_DEVICE_FOR_1D( n, i,
    {
        f(dst_data, src_data, src_start+i, dst_start+i);
    });
 
    Gpu::streamSynchronize();
}
 
template <typename DstTile1, typename DstTile2, typename SrcTile, typename F>
void transformParticles (DstTile1& dst1, DstTile2& dst2, const SrcTile& src, F&& f) noexcept
{
    auto np = src.numParticles();
    transformParticles(dst1, dst2, src, 0, 0, 0, np, std::forward<F>(f));
}
 
template <typename DstTile1, typename DstTile2, typename SrcTile,
          typename Index, typename N, typename F,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
void transformParticles (DstTile1& dst1, DstTile2& dst2, const SrcTile& src,
                         Index src_start, Index dst1_start, Index dst2_start, N n, F const& f) noexcept
{
    const auto src_data  = src.getConstParticleTileData();
          auto dst1_data = dst1.getParticleTileData();
          auto dst2_data = dst2.getParticleTileData();
 
    AMREX_HOST_DEVICE_FOR_1D( n, i,
    {
        f(dst1_data, dst2_data, src_data, src_start+i, dst1_start+i, dst2_start+i);
    });
 
    Gpu::streamSynchronize();
}
 
template <typename DstTile, typename SrcTile, typename Index, typename N,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
Index filterParticles (DstTile& dst, const SrcTile& src, const Index* mask) noexcept
{
    return filterParticles(dst, src, mask, 0, 0, src.numParticles());
}
 
template <typename DstTile, typename SrcTile, typename Index, typename N,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
Index filterParticles (DstTile& dst, const SrcTile& src, const Index* mask,
                       Index src_start, Index dst_start, N n) noexcept
{
    Gpu::DeviceVector<Index> offsets(n);
    Gpu::exclusive_scan(mask, mask+n, offsets.begin());
 
    Index last_mask=0, last_offset=0;
    Gpu::copyAsync(Gpu::deviceToHost, mask+n-1, mask + n, &last_mask);
    Gpu::copyAsync(Gpu::deviceToHost, offsets.data()+n-1, offsets.data()+n, &last_offset);
 
    auto* p_offsets = offsets.dataPtr();
 
    const auto src_data = src.getConstParticleTileData();
          auto dst_data = dst.getParticleTileData();
 
    AMREX_HOST_DEVICE_FOR_1D( n, i,
    {
        if (mask[i]) { copyParticle(dst_data, src_data, src_start+i, dst_start+p_offsets[i]); }
    });
 
    Gpu::streamSynchronize();
    return last_mask + last_offset;
}
 
template <typename DstTile, typename SrcTile, typename Pred,
          std::enable_if_t<!std::is_pointer_v<std::decay_t<Pred>>,int> foo = 0>
int filterParticles (DstTile& dst, const SrcTile& src, Pred&& p) noexcept
{
    return filterParticles(dst, src, std::forward<Pred>(p), 0, 0, src.numParticles());
}
 
template <typename DstTile, typename SrcTile, typename Pred, typename Index, typename N,
          std::enable_if_t<!std::is_pointer_v<std::decay_t<Pred>>,Index> nvccfoo = 0>
Index filterParticles (DstTile& dst, const SrcTile& src, Pred const& p,
                       Index src_start, Index dst_start, N n) noexcept
{
    Gpu::DeviceVector<Index> mask(n);
 
    auto* p_mask = mask.dataPtr();
    const auto src_data = src.getConstParticleTileData();
 
    amrex::ParallelForRNG(n,
    [p, p_mask, src_data, src_start] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept
    {
        amrex::ignore_unused(p, p_mask, src_data, src_start, engine);
        if constexpr (IsCallable<Pred,decltype(src_data),Index,RandomEngine>::value) {
            p_mask[i] = p(src_data, src_start+i, engine);
        } else {
            p_mask[i] = p(src_data, src_start+i);
        }
    });
    return filterParticles(dst, src, mask.dataPtr(), src_start, dst_start, n);
}
 
template <typename DstTile, typename SrcTile, typename Index, typename F,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
Index filterAndTransformParticles (DstTile& dst, const SrcTile& src, Index* mask, F const& f,
                                   Index src_start, Index dst_start) noexcept
{
    auto np = src.numParticles();
    Gpu::DeviceVector<Index> offsets(np);
    Gpu::exclusive_scan(mask, mask+np, offsets.begin());
 
    Index last_mask=0, last_offset=0;
    Gpu::copyAsync(Gpu::deviceToHost, mask+np-1, mask + np, &last_mask);
    Gpu::copyAsync(Gpu::deviceToHost, offsets.data()+np-1, offsets.data()+np, &last_offset);
 
    auto const* p_offsets = offsets.dataPtr();
 
    const auto src_data = src.getConstParticleTileData();
          auto dst_data = dst.getParticleTileData();
 
    AMREX_HOST_DEVICE_FOR_1D( np, i,
    {
        if (mask[i]) {
            f(dst_data, src_data, src_start+i,
              dst_start+p_offsets[src_start+i]);
        }
    });
 
    Gpu::streamSynchronize();
    return last_mask + last_offset;
}
 
template <typename DstTile, typename SrcTile, typename Index, typename F,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
Index filterAndTransformParticles (DstTile& dst, const SrcTile& src, Index* mask, F&& f) noexcept
{
    return filterAndTransformParticles(dst, src, mask, std::forward<F>(f), 0, 0);
}
 
template <typename DstTile, typename SrcTile, typename Pred, typename F,
          std::enable_if_t<!std::is_pointer_v<std::decay_t<Pred>>,int> foo = 0>
int filterAndTransformParticles (DstTile& dst, const SrcTile& src, Pred&& p, F&& f) noexcept
{
    return filterAndTransformParticles(dst, src, std::forward<Pred>(p), std::forward<F>(f), 0, 0);
}
 
template <typename DstTile1, typename DstTile2, typename SrcTile, typename Index, typename F,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
Index filterAndTransformParticles (DstTile1& dst1, DstTile2& dst2,
                                   const SrcTile& src, Index* mask, F const& f) noexcept
{
    auto np = src.numParticles();
    Gpu::DeviceVector<Index> offsets(np);
    Gpu::exclusive_scan(mask, mask+np, offsets.begin());
 
    Index last_mask=0, last_offset=0;
    Gpu::copyAsync(Gpu::deviceToHost, mask+np-1, mask + np, &last_mask);
    Gpu::copyAsync(Gpu::deviceToHost, offsets.data()+np-1, offsets.data()+np, &last_offset);
 
    auto* p_offsets = offsets.dataPtr();
 
    const auto src_data  = src.getConstParticleTileData();
          auto dst_data1 = dst1.getParticleTileData();
          auto dst_data2 = dst2.getParticleTileData();
 
    AMREX_HOST_DEVICE_FOR_1D( np, i,
    {
        if (mask[i]) { f(dst_data1, dst_data2, src_data, i, p_offsets[i], p_offsets[i]); }
    });
 
    Gpu::streamSynchronize();
    return last_mask + last_offset;
}
 
template <typename DstTile1, typename DstTile2, typename SrcTile, typename Pred, typename F,
          std::enable_if_t<!std::is_pointer_v<std::decay_t<Pred>>, int> foo = 0>
int filterAndTransformParticles (DstTile1& dst1, DstTile2& dst2, const SrcTile& src,
                                 Pred const& p, F&& f) noexcept
{
    auto np = src.numParticles();
    Gpu::DeviceVector<int> mask(np);
 
    auto* p_mask = mask.dataPtr();
    const auto src_data = src.getConstParticleTileData();
 
    amrex::ParallelForRNG(np,
    [p, p_mask, src_data] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept
    {
        amrex::ignore_unused(p, p_mask, src_data, engine);
        if constexpr (IsCallable<Pred,decltype(src_data),int,RandomEngine>::value) {
            p_mask[i] = p(src_data, i, engine);
        } else {
            p_mask[i] = p(src_data, i);
        }
    });
    return filterAndTransformParticles(dst1, dst2, src, mask.dataPtr(), std::forward<F>(f));
}
 
 
template <typename DstTile, typename SrcTile, typename Pred, typename F, typename Index,
          std::enable_if_t<!std::is_pointer_v<std::decay_t<Pred>>,Index> nvccfoo = 0>
Index filterAndTransformParticles (DstTile& dst, const SrcTile& src, Pred const& p, F&& f,
                                   Index src_start, Index dst_start) noexcept
{
    auto np = src.numParticles();
    Gpu::DeviceVector<Index> mask(np);
 
    auto* p_mask = mask.dataPtr();
    const auto src_data = src.getConstParticleTileData();
 
    amrex::ParallelForRNG(np,
    [p, p_mask, src_data, src_start] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept
    {
        amrex::ignore_unused(p, p_mask, src_data, src_start, engine);
        if constexpr (IsCallable<Pred,decltype(src_data),Index,RandomEngine>::value) {
            p_mask[i] = p(src_data, src_start+i, engine);
        } else {
            p_mask[i] = p(src_data, src_start+i);
        }
    });
    return filterAndTransformParticles(dst, src, mask.dataPtr(), std::forward<F>(f), src_start, dst_start);
}
 
 
template <typename PTile, typename N, typename Index,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
void gatherParticles (PTile& dst, const PTile& src, N np, const Index* inds)
{
    const auto src_data = src.getConstParticleTileData();
          auto dst_data = dst.getParticleTileData();
 
    AMREX_HOST_DEVICE_FOR_1D( np, i,
    {
        copyParticle(dst_data, src_data, inds[i], i);
    });
 
    Gpu::streamSynchronize();
}
 
template <typename PTile, typename N, typename Index,
          std::enable_if_t<std::is_integral_v<Index>, int> foo = 0>
void scatterParticles (PTile& dst, const PTile& src, N np, const Index* inds)
{
    const auto src_data = src.getConstParticleTileData();
          auto dst_data = dst.getParticleTileData();
 
    AMREX_HOST_DEVICE_FOR_1D( np, i,
    {
        copyParticle(dst_data, src_data, i, inds[i]);
    });
 
    Gpu::streamSynchronize();
}
 
}
 
#endif // include guard