Add TypeMultiplier, MakeZeroTuple and IdentityTuple

[AlexanderSinn]

Summary

This PR adds TypeMultiplier and MakeZeroTuple from Hi-PACE/hipace#1052 to AMReX.

TypeMultiplier can be used to shorten ReduceOps and ReduceData definitions where the same type is used many times.

MakeZeroTuple can be used to initialize a GpuTuple to zero (the default constructor would leave the values uninitialized).

IdentityTuple can be used to initialize a GpuTuple to the identity elements of each operation in a ReduceOps.

Additional background

Checklist

The proposed changes:

• fix a bug or incorrect behavior in AMReX
• add new capabilities to AMReX
• changes answers in the test suite to more than roundoff level
• are likely to significantly affect the results of downstream AMReX users
• include documentation in the code and/or rst files, if appropriate

[WeiqunZhang]

Any ideas on having something similar to MakeZeroTuple but suitable for reduce min and max also, like amrex::Reduce::detail::for_each_init?

[AlexanderSinn]

That would be a better version of MakeZeroTuple. I did not make it that way because the reductions in HiPACE++ don’t need it. It does seem a bit difficult to get the ReduceOps template list into a device lambda in user code, without capturing the ReduceOps object. I think an additional object such as a TypeList is necessary to store all the ReduceOpMax, ReduceOpSum etc so they don’t have the be manually given to the Tuple init function.

[AlexanderSinn]

Actually, maybe it would be fine to capture the ReduceOps object in a device lambda.

[WeiqunZhang]

Could something like this work? (Not tested).

template <typename... OPs, typename... Ts>
constexpr typename ReduceData<Ts...>::Type
MakeDefaultReduceTuple (TypeList<Ops...> op_list, TypeList<Ts...> type_list)
{   
    using ReduceTuple = typename ReduceData<Ts...>::Type;
    ReduceTuple r;
    return Reduce::detail::for_each_init<0, ReduceTuple, Ops...>(r);
}

[WeiqunZhang]

I guess it cannot be constexpr because for_each_init is not.

[WeiqunZhang]

But we probably could make for_each_init constexpr.

[AlexanderSinn]

I made the function use the GpuTuple and ReduceOps to make it more convenient in a typical reduce op use case. The following code now compiles for CPU and CUDA:

        amrex::TypeMultiplier<amrex::ReduceOps, amrex::ReduceOpSum[m_insitu_nrp + m_insitu_nip]> reduce_op;
        amrex::TypeMultiplier<amrex::ReduceData, amrex::Real[m_insitu_nrp], int[m_insitu_nip]> reduce_data(reduce_op);
        using ReduceTuple = typename decltype(reduce_data)::Type;
        reduce_op.eval(
            num_particles, reduce_data,
            [=] AMREX_GPU_DEVICE (int ip) -> ReduceTuple
            {
                const amrex::Real x = ptd.pos(0, ip);
                const amrex::Real y = ptd.pos(1, ip);
                const amrex::Real ux = ptd.rdata(PlasmaIdx::ux)[ip] * clight_inv; // proper velocity to u
                const amrex::Real uy = ptd.rdata(PlasmaIdx::uy)[ip] * clight_inv;
                const amrex::Real psi = ptd.rdata(PlasmaIdx::psi)[ip];

                if (ptd.id(ip) < 0 || x*x + y*y > insitu_radius_sq) {
                    return amrex::IdentityTuple(ReduceTuple{}, reduce_op);
                }
                // particle's Lorentz factor
                const amrex::Real gamma = (1.0_rt + ux*ux + uy*uy + psi*psi)/(2.0_rt*psi);
                // the *c from uz cancels with the /c from the proper velocity conversion
                const amrex::Real uz = (gamma - psi);
                // weight with quasi-static weighting factor
                const amrex::Real w = ptd.rdata(PlasmaIdx::w)[ip] * gamma/psi;
                // no quasi-static weighting factor to calculate quasi-static energy
                const amrex::Real energy = ptd.rdata(PlasmaIdx::w)[ip] * (gamma - 1._rt);
                return {            // Tuple contains:
                    w,              // 0    sum(w)
                    w*x,            // 1    [x]
                    w*x*x,          // 2    [x^2]
                    w*y,            // 3    [y]
                    w*y*y,          // 4    [y^2]
                    w*ux,           // 5    [ux]
                    w*ux*ux,        // 6    [ux^2]
                    w*uy,           // 7    [uy]
                    w*uy*uy,        // 8    [uy^2]
                    w*uz,           // 9    [uz]
                    w*uz*uz,        // 10   [uz^2]
                    w*gamma,        // 11   [ga]
                    w*gamma*gamma,  // 12   [ga^2]
                    energy,         // 13   [(ga-1)*(1-vz)]
                    1               // 14   Np
                };
            });

[WeiqunZhang]

LGTM