ParallelReduceTools.h

Go to the documentation of this file.

#ifndef PARALLEL_REDUCE_TOOLS_H
#define PARALLEL_REDUCE_TOOLS_H
 
#include <utility>  // for std::index_sequence
#include <variant>  // for std::variant
 
namespace ParticleBinning {
 
    // Possibility to manually choose between different reduction methods!
    enum class HistoReductionMode {
        Standard,        // Default auto selection
        ParallelReduce,  // Force usage of parallel_reduce if binCount <= maxArrSize
        TeamBased,       // Force team-based/atomic reduction if gpu enabled
        HostOnly
    };
 
    template <typename SizeType, typename IndexType, IndexType N>
    struct ArrayReduction {
        SizeType the_array[N];
 
        KOKKOS_INLINE_FUNCTION
        ArrayReduction() {
            for (IndexType i = 0; i < N; i++) {
                the_array[i] = 0;
            }
        }
        KOKKOS_INLINE_FUNCTION
        ArrayReduction(const ArrayReduction& rhs) {
            for (IndexType i = 0; i < N; i++) {
                the_array[i] = rhs.the_array[i];
            }
        }
        KOKKOS_INLINE_FUNCTION
        ArrayReduction& operator=(const ArrayReduction& rhs) {
            if (this != &rhs) {
                for (IndexType i = 0; i < N; ++i) {
                    the_array[i] = rhs.the_array[i];
                }
            }
            return *this;
        }
        KOKKOS_INLINE_FUNCTION
        ArrayReduction& operator+=(const ArrayReduction& src) {
            for (IndexType i = 0; i < N; i++) {
                the_array[i] += src.the_array[i];
            }
            return *this;
        }
    };
 
    /*
    Define logic for maxArrSize different reducer array types where N \in [1, ..., maxArrSize]
    */
 
    template <typename IndexType>
    constexpr IndexType maxArrSize = 5;
 
    template <typename SizeType, typename IndexType, typename Sequence>
    struct ReductionVariantHelper;
 
    template <typename SizeType, typename IndexType, IndexType... Sizes>
    struct ReductionVariantHelper<SizeType, IndexType, std::integer_sequence<IndexType, Sizes...>> {
        using type = std::variant<ArrayReduction<SizeType, IndexType, Sizes + 1>...>;
    };
 
    template <typename SizeType, typename IndexType>
    using ReductionVariant = typename ReductionVariantHelper<
            SizeType, IndexType,
            std::make_integer_sequence<IndexType, maxArrSize<IndexType>>>::type;
 
    template <typename SizeType, typename IndexType, IndexType N>
    ReductionVariant<SizeType, IndexType> createReductionObjectHelper(IndexType binCount) {
        if constexpr (N > maxArrSize<IndexType>) {
            throw std::out_of_range("binCount is out of the allowed range");
        } else if (binCount == N) {
            return ArrayReduction<SizeType, IndexType, N>();
        } else {
            return createReductionObjectHelper<SizeType, IndexType, N + 1>(binCount);
        }
    }
 
    template <typename SizeType, typename IndexType>
    ReductionVariant<SizeType, IndexType> createReductionObject(IndexType binCount) {
        return createReductionObjectHelper<SizeType, IndexType, 1>(binCount);
    }
 
    template <typename SizeType, typename IndexType>
    struct HostArrayReduction {
        SizeType* the_array;
 
        static IndexType binCountStatic;
 
#ifndef OPALX_DEVICE_COMPILATION
        HostArrayReduction() {
            the_array = new SizeType[binCountStatic];
            for (IndexType i = 0; i < binCountStatic; i++) {
                the_array[i] = 0;
            }
        }
 
        HostArrayReduction(const HostArrayReduction& rhs) {
            the_array = new SizeType[binCountStatic];
            for (IndexType i = 0; i < binCountStatic; i++) {
                the_array[i] = rhs.the_array[i];
            }
        }
 
        ~HostArrayReduction() { delete[] the_array; }
 
        HostArrayReduction& operator=(const HostArrayReduction& rhs) {
            the_array = new SizeType[binCountStatic];
            if (this != &rhs) {
                for (IndexType i = 0; i < binCountStatic; ++i) {
                    the_array[i] = rhs.the_array[i];
                }
            }
            return *this;
        }
 
        HostArrayReduction& operator+=(const HostArrayReduction& src) {
            for (IndexType i = 0; i < binCountStatic; i++) {
                the_array[i] += src.the_array[i];
            }
            return *this;
        }
#else
        KOKKOS_INLINE_FUNCTION
        HostArrayReduction& operator+=(const HostArrayReduction& src) {
            (void)src;  // Silence unused parameter warning on Clang
            Kokkos::abort(
                    "Error: HostArrayReduction is not supported on device backends "
                    "(CUDA/HIP/SYCL)!\n       Note: It exists only for compilation compatibility.");
            return *this;
        }
#endif
    };
 
    template <typename SizeType, typename IndexType>
    IndexType HostArrayReduction<SizeType, IndexType>::binCountStatic = 10;
 
    template <typename T, unsigned Dim>
    T vnorm(const VField_t<T, Dim>& field, int p = 2) {
        T sum = 0;
        ippl::parallel_reduce(
                "VectorFieldNormReduce", field.getFieldRangePolicy(),
                KOKKOS_LAMBDA(const ippl::RangePolicy<Dim>::index_array_type& idx, T& loc_sum) {
                    ippl::Vector<T, Dim> e = apply(field, idx);
                    loc_sum += std::pow(e.dot(e), p / 2.0);
                },
                Kokkos::Sum<T>(sum));
        return std::pow(sum, 1.0 / p);
    }
 
}  // namespace ParticleBinning
 
namespace Kokkos {
    template <typename SizeType, typename IndexType, IndexType N>
    struct reduction_identity<ParticleBinning::ArrayReduction<SizeType, IndexType, N>> {
        KOKKOS_FORCEINLINE_FUNCTION static ParticleBinning::ArrayReduction<SizeType, IndexType, N>
        sum() {
            return ParticleBinning::ArrayReduction<SizeType, IndexType, N>();
        }
    };
 
    template <typename SizeType, typename IndexType>
    struct reduction_identity<ParticleBinning::HostArrayReduction<SizeType, IndexType>> {
        KOKKOS_FORCEINLINE_FUNCTION static ParticleBinning::HostArrayReduction<SizeType, IndexType>
        sum() {
            return ParticleBinning::HostArrayReduction<SizeType, IndexType>();
        }
    };
}  // namespace Kokkos
 
#endif  // PARALLEL_REDUCE_TOOLS_H