beamoptic.H

Go to the documentation of this file.

/* Copyright 2022-2023 The Regents of the University of California, through Lawrence
 *           Berkeley National Laboratory (subject to receipt of any required
 *           approvals from the U.S. Dept. of Energy). All rights reserved.
 *
 * This file is part of ImpactX.
 *
 * Authors: Axel Huebl
 * License: BSD-3-Clause-LBNL
 */
#ifndef IMPACTX_ELEMENTS_MIXIN_BEAMOPTIC_H
#define IMPACTX_ELEMENTS_MIXIN_BEAMOPTIC_H
 
#include "particles/ImpactXParticleContainer.H"
#include "particles/PushAll.H"
 
#include <AMReX_Extension.H> // for AMREX_RESTRICT
#include <AMReX_REAL.H>
#include <AMReX_SIMD.H>
 
#include <type_traits>
 
 
namespace impactx::elements::mixin
{
namespace detail
{
    template <typename T_Element>
    struct PushSingleParticle
    {
        using PType = ImpactXParticleContainer::ParticleType;
        using ParticleTileType = amrex::ParticleTile<amrex::SoAParticle<RealSoA::nattribs, IntSoA::nattribs>,RealSoA::nattribs, IntSoA::nattribs>;

        PushSingleParticle (T_Element element,
                            amrex::ParticleReal* AMREX_RESTRICT part_x,
                            amrex::ParticleReal* AMREX_RESTRICT part_y,
                            amrex::ParticleReal* AMREX_RESTRICT part_t,
                            amrex::ParticleReal* AMREX_RESTRICT part_px,
                            amrex::ParticleReal* AMREX_RESTRICT part_py,
                            amrex::ParticleReal* AMREX_RESTRICT part_pt,
                            uint64_t* AMREX_RESTRICT part_idcpu,
                            RefPart ref_part)
            : m_element(std::move(element)),
              m_part_x(part_x), m_part_y(part_y), m_part_t(part_t),
              m_part_px(part_px), m_part_py(part_py), m_part_pt(part_pt),
              m_part_idcpu(part_idcpu),
              m_ref_part(std::move(ref_part))
        {
            // pre-compute and cache constants that are independent of the
            // individually tracked particle
            m_element.compute_constants(m_ref_part);
        }
 
        PushSingleParticle () = delete;
        PushSingleParticle (PushSingleParticle const &) = default;
        PushSingleParticle (PushSingleParticle &&) = default;
        ~PushSingleParticle () = default;

        AMREX_GPU_DEVICE AMREX_FORCE_INLINE
        void
        operator() (int i) const
        {
            // access SoA Real data
            // note: an optimizing compiler will eliminate loads of unused parameters
            amrex::ParticleReal & AMREX_RESTRICT x = m_part_x[i];
            amrex::ParticleReal & AMREX_RESTRICT y = m_part_y[i];
            amrex::ParticleReal & AMREX_RESTRICT t = m_part_t[i];
            amrex::ParticleReal & AMREX_RESTRICT px = m_part_px[i];
            amrex::ParticleReal & AMREX_RESTRICT py = m_part_py[i];
            amrex::ParticleReal & AMREX_RESTRICT pt = m_part_pt[i];
            uint64_t & AMREX_RESTRICT idcpu = m_part_idcpu[i];
 
            // push through element
            m_element(x, y, t, px, py, pt, idcpu, m_ref_part);
        }
 
#ifdef AMREX_USE_SIMD
        template<int SIMD_WIDTH>
        AMREX_GPU_DEVICE AMREX_FORCE_INLINE
        void
        operator() (amrex::SIMDindex<SIMD_WIDTH, int> idx) const
        {
            using namespace amrex::simd;
 
            int const i = idx.index;
 
            // uninitialized memory
            SIMDParticleReal<SIMD_WIDTH> part_x;
            SIMDParticleReal<SIMD_WIDTH> part_y;
            SIMDParticleReal<SIMD_WIDTH> part_t;
            SIMDParticleReal<SIMD_WIDTH> part_px;
            SIMDParticleReal<SIMD_WIDTH> part_py;
            SIMDParticleReal<SIMD_WIDTH> part_pt;
            SIMDIdCpu<decltype(part_x)> part_idcpu;
 
            // initialize vector registers
            // TODO stdx::vector_aligned needs alignment guarantees
            //      https://github.com/AMReX-Codes/amrex/issues/4592
            //      https://en.cppreference.com/w/cpp/experimental/simd/simd/copy_from
            part_x.copy_from(&m_part_x[i], stdx::element_aligned);
            part_y.copy_from(&m_part_y[i], stdx::element_aligned);
            part_t.copy_from(&m_part_t[i], stdx::element_aligned);
            part_px.copy_from(&m_part_px[i], stdx::element_aligned);
            part_py.copy_from(&m_part_py[i], stdx::element_aligned);
            part_pt.copy_from(&m_part_pt[i], stdx::element_aligned);
            part_idcpu.copy_from(&m_part_idcpu[i], stdx::element_aligned);
 
            // push through element
            m_element(part_x, part_y, part_t, part_px, part_py, part_pt, part_idcpu, m_ref_part);
 
            // write back to memory
            // TODO stdx::vector_aligned needs alignment guarantees
            //      https://github.com/AMReX-Codes/amrex/issues/4592
            //      https://en.cppreference.com/w/cpp/experimental/simd/simd/copy_from
            if constexpr (is_nth_arg_non_const(&T_Element::template operator()<decltype(part_x), decltype(part_idcpu)>, 0))
                part_x.copy_to(&m_part_x[i], stdx::element_aligned);
            if constexpr (is_nth_arg_non_const(&T_Element::template operator()<decltype(part_x), decltype(part_idcpu)>, 1))
                part_y.copy_to(&m_part_y[i], stdx::element_aligned);
            if constexpr (is_nth_arg_non_const(&T_Element::template operator()<decltype(part_x), decltype(part_idcpu)>, 2))
                part_t.copy_to(&m_part_t[i], stdx::element_aligned);
            if constexpr (is_nth_arg_non_const(&T_Element::template operator()<decltype(part_x), decltype(part_idcpu)>, 3))
                part_px.copy_to(&m_part_px[i], stdx::element_aligned);
            if constexpr (is_nth_arg_non_const(&T_Element::template operator()<decltype(part_x), decltype(part_idcpu)>, 4))
                part_py.copy_to(&m_part_py[i], stdx::element_aligned);
            if constexpr (is_nth_arg_non_const(&T_Element::template operator()<decltype(part_x), decltype(part_idcpu)>, 5))
                part_pt.copy_to(&m_part_pt[i], stdx::element_aligned);
            if constexpr (is_nth_arg_non_const(&T_Element::template operator()<decltype(part_x), decltype(part_idcpu)>, 6))
                part_idcpu.copy_to(&m_part_idcpu[i], stdx::element_aligned);
        }
#endif
 
    private:
        T_Element m_element;
        amrex::ParticleReal* const AMREX_RESTRICT m_part_x;
        amrex::ParticleReal* const AMREX_RESTRICT m_part_y;
        amrex::ParticleReal* const AMREX_RESTRICT m_part_t;
        amrex::ParticleReal* const AMREX_RESTRICT m_part_px;
        amrex::ParticleReal* const AMREX_RESTRICT m_part_py;
        amrex::ParticleReal* const AMREX_RESTRICT m_part_pt;
        uint64_t* const AMREX_RESTRICT m_part_idcpu;
        RefPart m_ref_part;
    };

    template< typename T_Element >
    void push_all_particles (
        ImpactXParticleContainer::iterator & pti,
        RefPart & AMREX_RESTRICT ref_part,
        T_Element & element
    ) {
        const int np = pti.numParticles();
 
        // preparing access to particle data: SoA of Reals
        auto& soa_real = pti.GetStructOfArrays().GetRealData();
        amrex::ParticleReal* const AMREX_RESTRICT part_x = soa_real[RealSoA::x].dataPtr();
        amrex::ParticleReal* const AMREX_RESTRICT part_y = soa_real[RealSoA::y].dataPtr();
        amrex::ParticleReal* const AMREX_RESTRICT part_t = soa_real[RealSoA::t].dataPtr();
        amrex::ParticleReal* const AMREX_RESTRICT part_px = soa_real[RealSoA::px].dataPtr();
        amrex::ParticleReal* const AMREX_RESTRICT part_py = soa_real[RealSoA::py].dataPtr();
        amrex::ParticleReal* const AMREX_RESTRICT part_pt = soa_real[RealSoA::pt].dataPtr();
 
        uint64_t* const AMREX_RESTRICT part_idcpu = pti.GetStructOfArrays().GetIdCPUData().dataPtr();
 
        detail::PushSingleParticle<T_Element> const pushSingleParticle(
            element, part_x, part_y, part_t, part_px, part_py, part_pt, part_idcpu, ref_part);
 
        // loop over beam particles in the box
#ifdef AMREX_USE_SIMD
        if constexpr (amrex::simd::is_vectorized<T_Element>) {
            amrex::ParallelForSIMD<T_Element::simd_width>(np, pushSingleParticle);  // TODO: test 2,4,8, ... and handle remainder
        } else
#endif
        {
            amrex::ParallelFor(np, pushSingleParticle);
        }
    }
} // namespace detail

    template<typename T_Element>
    struct BeamOptic
    {
        void operator() (
            ImpactXParticleContainer & pc,
            int step,
            int period
        )
        {
            static_assert(
                std::is_base_of_v<BeamOptic, T_Element>,
                "BeamOptic can only be used as a mixin class!"
            );
 
            T_Element& element = *static_cast<T_Element*>(this);
            push_all(pc, element, step, period);
        }

         void operator() (
            ImpactXParticleContainer::iterator & pti,
            RefPart & AMREX_RESTRICT ref_part
         )
         {
            static_assert(
                std::is_base_of_v<BeamOptic, T_Element>,
                "BeamOptic can only be used as a mixin class!"
            );
 
            T_Element& element = *static_cast<T_Element*>(this);
            detail::push_all_particles<T_Element>(pti, ref_part, element);
         }
    };
 
} // namespace impactx::elements::mixin
 
#endif // IMPACTX_ELEMENTS_MIXIN_BEAMOPTIC_H