How to generate nested loops at compile time

亡梦爱人 提交于 2021-02-17 09:14:37

问题


I have an integer N which I know at compile time. I also have an std::array holding integers describing the shape of an N-dimensional array. I want to generate nested loops, as described bellow, at compile time, using metaprogramming techniques.

constexpr int N {4};
constexpr std::array<int, N> shape {{1,3,5,2}};


auto f = [/* accept object which uses coords */] (auto... coords) { 
     // do sth with coords
}; 

// This is what I want to generate.
for(int i = 0; i < shape[0]; i++) {
     for(int j = 0; j < shape[1]; j++) {
          for(int k = 0; k < shape[2]; k++) {
                for(int l = 0; l < shape[3]; l++) {
                    f(i,j,k,l) // object is modified via the lambda function.
                }
          }
     }
}

Note the parameter N is known at compile time but might change unpredictably between compilations, hence I can't hard code the loops as above. Ideally the loop generation mechanism will provide an interface which accepts the lambda function, generates the loops and calls the function producing the equivalent code as above. I am aware that one can write an equivalent loop at runtime with a single while loop and an array of indices, and there are answers to this question already. I am, however, not interested in this solution. I am also not interested in solutions involving preprocessor magic.


回答1:


Something like this (NOTE: I take the "shape" as a variadic template argument set..)

#include <iostream>

template <int I, int ...N>
struct Looper{
    template <typename F, typename ...X>
    constexpr void operator()(F& f, X... x) {
        for (int i = 0; i < I; ++i) {
            Looper<N...>()(f, x..., i);
        }
    }
};

template <int I>
struct Looper<I>{
    template <typename F, typename ...X>
    constexpr void operator()(F& f, X... x) {
        for (int i = 0; i < I; ++i) {
            f(x..., i);
        }
    }
};

int main()
{
    int v = 0;
    auto f = [&](int i, int j, int k, int l) {
        v += i + j + k + l;
    };

    Looper<1, 3, 5, 2>()(f);

    auto g = [&](int i) {
        v += i;
    };

    Looper<5>()(g);

    std::cout << v << std::endl;
}



回答2:


Assuming you don't want total loop unrolling, just generation of i, j, k etc. argument tuples for f:

#include <stdio.h>
#include <utility>      // std::integer_sequence

template< int dim >
constexpr auto item_size_at()
    -> int
{ return ::shape[dim + 1]*item_size_at<dim + 1>(); }

template<> constexpr auto item_size_at<::N-1>() -> int { return 1; }

template< size_t... dim >
void call_f( int i, std::index_sequence<dim...> )
{
    f( (i/item_size_at<dim>() % ::shape[dim])... );
}

auto main()
    -> int
{
    int const n_items = ::shape[0]*item_size_at<0>();
    for( int i = 0; i < n_items; ++i )
    {
        call_f( i, std::make_index_sequence<::N>() );
    }
}



回答3:


I suppose this is exactly what you asked for:

#include <array>
#include <iostream>

constexpr int N{4};
constexpr std::array<int, N> shape {{1,3,5,2}};

// Diagnositcs

template<typename V, typename ...Vals>
struct TPrintf {
        constexpr static void call(V v, Vals ...vals) {
                std::cout << v << " ";
                TPrintf<Vals...>::call(vals...);
        }
};

template<typename V>
struct TPrintf<V> {
        constexpr static void call(V v) {
                std::cout << v << std::endl;
        }
};


template<typename ...Vals>
constexpr void t_printf(Vals ...vals) {
        TPrintf<Vals...>::call(vals...);
}

// Unroll

template<int CtIdx, typename F>
struct NestedLoops {
        template<typename ...RtIdx>
        constexpr static void call(const F& f, RtIdx ...idx) {
                for(int i = 0; i < shape[CtIdx]; ++i) {
                        NestedLoops<CtIdx + 1, F>::call(f, idx..., i);
                }
        }
};

template<typename F>
struct NestedLoops<N-1, F> {
        template<typename ...RtIdx>
        constexpr static void call(const F& f, RtIdx ...idx) {
                for(int i = 0; i < shape[N-1]; ++i) {
                        f(idx..., i);
                }
        }
};

template<typename F>
void nested_loops(const F& f) {
        NestedLoops<0, F>::call(f);
}

int main()
{
        auto lf = [](int i, int j, int k, int l) {
                t_printf(i,j,k,l);
        };

        nested_loops(lf);
        return 0;
}



回答4:


Another variant of the same thing:

template <size_t shape_index, size_t shape_size>
struct Looper
{
    template <typename Functor>
    void operator()(const std::array<int, shape_size>& shape, Functor functor)
    {
        for (int index = 0; index < shape[shape_index]; ++index)
        {
            Looper<shape_index + 1, shape_size>()
                (
                    shape,
                    [index, &functor](auto... tail){ functor(index, tail...); }
                );
        }
    }
};

template <size_t shape_size>
struct Looper<shape_size, shape_size>
{
    template <typename Functor>
    void operator()(const std::array<int, shape_size>&, Functor functor)
    {
        functor();
    }
};

template <size_t shape_size, typename Functor>
void loop(const std::array<int, shape_size>& shape, Functor functor)
{
    Looper<0, shape_size>()(shape, functor);
}

Example of use:

constexpr size_t N {4};

constexpr std::array<int, N> shape {{1,3,5,2}};

void f(int i, int j, int k, int l)
{
    std::cout
        << std::setw(5) << i
        << std::setw(5) << j
        << std::setw(5) << k
        << std::setw(5) << l
        << std::endl;
}

// ...

loop(shape, f);

Live demo



来源:https://stackoverflow.com/questions/38393694/how-to-generate-nested-loops-at-compile-time

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!