IdxRange.hpp

Go to the documentation of this file.

/* Copyright 2024 René Widera
 * SPDX-License-Identifier: MPL-2.0
 */
 
#pragma once
 
#include "alpaka/mem/IdxRange.hpp"
#include "alpaka/onAcc/WorkerGroup.hpp"
#include "alpaka/onAcc/tag.hpp"
 
namespace alpaka::onAcc::internal
{
    template<typename T_ExtentFn>
    struct IdxRangeFn
    {
        constexpr IdxRangeFn(T_ExtentFn const& extentFn) : m_extentFn{extentFn}
        {
        }
 
        constexpr auto getIdxRange(auto const& acc) const
        {
            return IdxRange{m_extentFn(acc)};
        }
 
    private:
        T_ExtentFn const m_extentFn;
    };
 
    template<concepts::Origin T_Origin, concepts::Unit T_Unit, typename T_MultiDimensional = MultiDimensional<true>>
    struct IdxRangeLazy
    {
        constexpr IdxRangeLazy(T_Origin const& origin, T_Unit const& unit, T_MultiDimensional = T_MultiDimensional{})
        {
            alpaka::unused(origin, unit);
        }
 
        constexpr auto getIdxRange(auto const& acc) const
        {
            auto const extent
                = internalCompute::GetExtentsOf::Op<ALPAKA_TYPEOF(acc), T_Origin, T_Unit>{}(acc, T_Origin{}, T_Unit{});
 
            if constexpr(T_MultiDimensional::value == false)
                return IdxRange{Vec{extent.product()}};
            else
                return IdxRange{extent};
        }
    };
} // namespace alpaka::onAcc::internal
 
namespace alpaka::trait
{
    template<concepts::SpecializationOf<onAcc::internal::IdxRangeLazy> T>
    struct IsLazyIndexRange<T> : std::true_type
    {
    };
 
    template<concepts::SpecializationOf<onAcc::internal::IdxRangeFn> T>
    struct IsLazyIndexRange<T> : std::true_type
    {
    };
} // namespace alpaka::trait