AoSoA.hpp

Go to the documentation of this file.

// Copyright 2022 Bernhard Manfred Gruber
// SPDX-License-Identifier: MPL-2.0
 
#pragma once
 
#include "Common.hpp"
 
#include <limits>
 
namespace llama::mapping
{
    LLAMA_EXPORT
    template<typename RecordDim, std::size_t VectorRegisterBits>
    inline constexpr std::size_t maxLanes = []() constexpr
    {
        auto max = std::numeric_limits<std::size_t>::max();
        forEachLeafCoord<RecordDim>(
            [&](auto rc)
            {
                using AttributeType = GetType<RecordDim, decltype(rc)>;
                max = std::min(max, VectorRegisterBits / (sizeof(AttributeType) * CHAR_BIT));
            });
        return max;
    }();
 
    LLAMA_EXPORT
    template<
        typename TArrayExtents,
        typename TRecordDim,
        typename TArrayExtents::value_type Lanes,
        FieldAlignment TFieldAlignment = FieldAlignment::Align,
        typename TLinearizeArrayIndexFunctor = LinearizeArrayIndexRight,
        template<typename> typename PermuteFields = PermuteFieldsInOrder>
    struct AoSoA : MappingBase<TArrayExtents, TRecordDim>
    {
    private:
        using Base = MappingBase<TArrayExtents, TRecordDim>;
        using size_type = typename Base::size_type;
 
    public:
        inline static constexpr typename TArrayExtents::value_type lanes = Lanes;
        inline static constexpr FieldAlignment fieldAlignment = TFieldAlignment;
        using LinearizeArrayIndexFunctor = TLinearizeArrayIndexFunctor;
        using Permuter = PermuteFields<FlatRecordDim<TRecordDim>>;
        inline static constexpr std::size_t blobCount = 1;
 
#if defined(__NVCC__) && __CUDACC_VER_MAJOR__ >= 12
        using Base::Base;
#else
        constexpr AoSoA() = default;
 
        LLAMA_FN_HOST_ACC_INLINE constexpr explicit AoSoA(TArrayExtents extents, TRecordDim = {}) : Base(extents)
        {
        }
#endif
 
        LLAMA_FN_HOST_ACC_INLINE constexpr auto blobSize(size_type) const -> size_type
        {
            const auto rs = static_cast<size_type>(
                flatSizeOf<typename Permuter::FlatRecordDim, fieldAlignment == FieldAlignment::Align>);
            return roundUpToMultiple(LinearizeArrayIndexFunctor{}.size(Base::extents()) * rs, Lanes * rs);
        }
 
        template<std::size_t... RecordCoords>
        LLAMA_FN_HOST_ACC_INLINE constexpr auto blobNrAndOffset(
            typename Base::ArrayIndex ai,
            RecordCoord<RecordCoords...> rc = {}) const -> NrAndOffset<size_type>
        {
            return blobNrAndOffset(LinearizeArrayIndexFunctor{}(ai, Base::extents()), rc);
        }
 
        // Exposed for aosoaCommonBlockCopy. Should be private ...
        template<std::size_t... RecordCoords>
        LLAMA_FN_HOST_ACC_INLINE constexpr auto blobNrAndOffset(
            size_type flatArrayIndex,
            RecordCoord<RecordCoords...> = {}) const -> NrAndOffset<size_type>
        {
            constexpr std::size_t flatFieldIndex =
#if defined(__NVCC__) && __CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ <= 6
                *& // mess with nvcc compiler state to workaround bug
#endif
                 Permuter::template permute<flatRecordCoord<TRecordDim, RecordCoord<RecordCoords...>>>;
            const auto blockIndex = flatArrayIndex / Lanes;
            const auto laneIndex = flatArrayIndex % Lanes;
            const auto offset
                = static_cast<size_type>(
                      flatSizeOf<typename Permuter::FlatRecordDim, fieldAlignment == FieldAlignment::Align> * Lanes)
                    * blockIndex
                + static_cast<size_type>(flatOffsetOf<
                                         typename Permuter::FlatRecordDim,
                                         flatFieldIndex,
                                         fieldAlignment == FieldAlignment::Align>)
                    * Lanes
                + static_cast<size_type>(sizeof(GetType<TRecordDim, RecordCoord<RecordCoords...>>)) * laneIndex;
            return {0, offset};
        }
    };
 
    LLAMA_EXPORT
    template<
        std::size_t Lanes,
        FieldAlignment Alignment = FieldAlignment::Align,
        typename LinearizeArrayIndexFunctor = LinearizeArrayIndexRight,
        template<typename> typename PermuteFields = PermuteFieldsInOrder>
    struct BindAoSoA
    {
        template<typename ArrayExtents, typename RecordDim>
        using fn = AoSoA<ArrayExtents, RecordDim, Lanes, Alignment, LinearizeArrayIndexFunctor, PermuteFields>;
    };
 
    LLAMA_EXPORT
    template<typename Mapping>
    inline constexpr bool isAoSoA = false;
 
    LLAMA_EXPORT
    template<
        typename AD,
        typename RD,
        typename AD::value_type L,
        FieldAlignment A,
        typename Lin,
        template<typename>
        typename Perm>
    inline constexpr bool isAoSoA<AoSoA<AD, RD, L, A, Lin, Perm>> = true;
} // namespace llama::mapping